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

nette · Oct 4, 2022 · 6509699 · 6509699
1 parent 9073c8a
commit 6509699
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
@@ -62,8 +62,7 @@ public function extractMethodBodies(string $className): array
 	{
 		$nodeFinder = new NodeFinder;
 		$classNode = $nodeFinder->findFirst($this->statements, function (Node $node) use ($className) {
-			return ($node instanceof Node\Stmt\Class_ || $node instanceof Node\Stmt\Trait_)
-				&& $node->namespacedName->toString() === $className;
+			return $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);