Extractor: Fixed extracting enum method body [Closes #115] (#116)

nette · Sep 19, 2022 · 7e1f0c2 · 7e1f0c2
1 parent ce44913
commit 7e1f0c2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/PhpGenerator/Extractor.php b/src/PhpGenerator/Extractor.php
@@ -66,8 +66,7 @@ public function extractMethodBodies(string $className): array
 		$nodeFinder = new NodeFinder;
 		$classNode = $nodeFinder->findFirst(
 			$this->statements,
-			fn(Node $node) => ($node instanceof Node\Stmt\Class_ || $node instanceof Node\Stmt\Trait_)
-				&& $node->namespacedName->toString() === $className,
+			fn(Node $node) => $node instanceof Node\Stmt\ClassLike && $node->namespacedName->toString() === $className,
 		);
 
 		$res = [];

diff --git a/tests/PhpGenerator/Extractor.getMethodBodies.phpt b/tests/PhpGenerator/Extractor.getMethodBodies.phpt
@@ -35,6 +35,17 @@ abstract class Another
 		echo 123;
 	}
 }
+
+enum Color
+{
+	case Red;
+	case Blue;
+
+	public function getName(): string
+	{
+		return $this->name;
+	}
+}
 ');
 
 $bodies = $extractor->extractMethodBodies('NS\Undefined');
@@ -45,3 +56,8 @@ Assert::same([
 	'bar1' => "\$a = 10;\necho 123;",
 	'bar2' => 'echo "hello";',
 ], $bodies);
+
+$bodies = $extractor->extractMethodBodies('NS\Color');
+Assert::same([
+	'getName' => "return \$this->name;",
+], $bodies);