ClassReflection: make getTraits recursive

iamrgroot · web-flow · commit 07d78b21b3a3 · 2021-06-17T09:58:35.000+02:00
diff --git a/src/Reflection/ClassReflection.php b/src/Reflection/ClassReflection.php
@@ -670,13 +670,36 @@ public function getInterfaces(): array
 	}
 
 	/**
-	 * @return \PHPStan\Reflection\ClassReflection[]
+	 * @return array<string, \PHPStan\Reflection\ClassReflection>
 	 */
-	public function getTraits(): array
+	public function getTraits(bool $recursive = false): array
 	{
-		return array_map(function (\ReflectionClass $trait): ClassReflection {
+		$traits = [];
+
+		if ($recursive) {
+			foreach ($this->collectTraits($this->getNativeReflection()) as $trait) {
+				$traits[$trait->getName()] = $trait;
+			}
+		} else {
+			$traits = $this->getNativeReflection()->getTraits();
+		}
+
+		$traits = array_map(function (\ReflectionClass $trait): ClassReflection {
 			return $this->reflectionProvider->getClass($trait->getName());
-		}, $this->getNativeReflection()->getTraits());
+		}, $traits);
+
+		if ($recursive) {
+			$parentClass = $this->getNativeReflection()->getParentClass();
+
+			if ($parentClass !== false) {
+				return array_merge(
+					$traits,
+					$this->reflectionProvider->getClass($parentClass->getName())->getTraits(true)
+				);
+			}
+		}
+
+		return $traits;
 	}
 
 	/**
diff --git a/tests/PHPStan/Reflection/ClassReflectionTest.php b/tests/PHPStan/Reflection/ClassReflectionTest.php
@@ -220,4 +220,100 @@ public function testDeprecatedConstantFromAnotherFile(): void
 		$this->assertTrue($constant->isDeprecated()->yes());
 	}
 
+	/**
+	 * @dataProvider dataNestedRecursiveTraits
+	 * @param class-string $className
+	 * @param array<class-string, class-string> $expected
+	 * @param bool $recursive
+	 */
+	public function testGetTraits(string $className, array $expected, bool $recursive): void
+	{
+		$reflectionProvider = $this->createBroker();
+
+		$this->assertSame(
+			array_map(
+				static function (ClassReflection $classReflection): string {
+					return $classReflection->getNativeReflection()->getName();
+				},
+				$reflectionProvider->getClass($className)->getTraits($recursive)
+			),
+			$expected
+		);
+	}
+
+	public function dataNestedRecursiveTraits(): array
+	{
+		return [
+			[
+				\NestedTraits\NoTrait::class,
+				[],
+				false,
+			],
+			[
+				\NestedTraits\NoTrait::class,
+				[],
+				true,
+			],
+			[
+				\NestedTraits\Foo::class,
+				[
+					\NestedTraits\FooTrait::class => \NestedTraits\FooTrait::class,
+				],
+				false,
+			],
+			[
+				\NestedTraits\Foo::class,
+				[
+					\NestedTraits\FooTrait::class => \NestedTraits\FooTrait::class,
+				],
+				true,
+			],
+			[
+				\NestedTraits\Bar::class,
+				[
+					\NestedTraits\BarTrait::class => \NestedTraits\BarTrait::class,
+				],
+				false,
+			],
+			[
+				\NestedTraits\Bar::class,
+				[
+					\NestedTraits\BarTrait::class => \NestedTraits\BarTrait::class,
+					\NestedTraits\FooTrait::class => \NestedTraits\FooTrait::class,
+				],
+				true,
+			],
+			[
+				\NestedTraits\Baz::class,
+				[
+					\NestedTraits\BazTrait::class => \NestedTraits\BazTrait::class,
+				],
+				false,
+			],
+			[
+				\NestedTraits\Baz::class,
+				[
+					\NestedTraits\BazTrait::class => \NestedTraits\BazTrait::class,
+					\NestedTraits\BarTrait::class => \NestedTraits\BarTrait::class,
+					\NestedTraits\FooTrait::class => \NestedTraits\FooTrait::class,
+				],
+				true,
+			],
+			[
+				\NestedTraits\BazChild::class,
+				[],
+				false,
+			],
+			[
+				\NestedTraits\BazChild::class,
+				[
+					\NestedTraits\BazTrait::class => \NestedTraits\BazTrait::class,
+					\NestedTraits\BarTrait::class => \NestedTraits\BarTrait::class,
+					\NestedTraits\FooTrait::class => \NestedTraits\FooTrait::class,
+				],
+				true,
+			],
+		];
+	}
+
 }
diff --git a/tests/PHPStan/Reflection/data/NestedTraits.php b/tests/PHPStan/Reflection/data/NestedTraits.php
@@ -0,0 +1,40 @@
+<?php
+
+namespace NestedTraits;
+
+trait FooTrait
+{
+}
+
+trait BarTrait
+{
+    use FooTrait;
+}
+
+trait BazTrait
+{
+    use BarTrait;
+}
+
+class NoTrait
+{
+}
+
+class Foo
+{
+    use FooTrait;
+}
+
+class Bar
+{
+    use BarTrait;
+}
+
+class Baz
+{
+    use BazTrait;
+}
+
+class BazChild extends Baz
+{
+}
