doctrine-dbal: union query fetch types (#296)

Co-authored-by: Markus Staab <[email protected]>

Author: staabm

src/DoctrineReflection/DoctrineReflection.php

@@ -18,14 +18,42 @@
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\UnionType;
 use staabm\PHPStanDba\QueryReflection\QueryReflection;
 use staabm\PHPStanDba\QueryReflection\QueryReflector;
 use Traversable;
 
 final class DoctrineReflection
 {
-    public function reduceResultType(MethodReflection $methodReflection, Type $resultRowType): ?Type
+    public function reduceResultType(MethodReflection $methodReflection, Type $resultType): ?Type
     {
+        if ($resultType instanceof UnionType) {
+            $resultTypes = [];
+
+            foreach ($resultType->getTypes() as $type) {
+                $rowType = $this->reduceResultType($methodReflection, $type);
+                if (null === $rowType) {
+                    return null;
+                }
+                $resultTypes[] = $rowType;
+            }
+
+            return TypeCombinator::union(...$resultTypes);
+        }
+
+        if ($resultType instanceof GenericObjectType) {
+            $genericTypes = $resultType->getTypes();
+
+            if (1 !== \count($genericTypes)) {
+                return null;
+            }
+
+            $resultRowType = $genericTypes[0];
+
+            return $this->reduceResultType($methodReflection, $resultRowType);
+        }
+
+        $resultRowType = $resultType;
         $usedMethod = strtolower($methodReflection->getName());
 
         switch ($usedMethod) {

src/Extensions/DoctrineConnectionFetchDynamicReturnTypeExtension.php

@@ -85,12 +85,12 @@ private function inferType(MethodReflection $methodReflection, Expr $queryExpr,
         $doctrineReflection = new DoctrineReflection();
 
         if (null === $paramsExpr) {
-            $queryString = $queryReflection->resolveQueryStrings($queryExpr, $scope);
+            $queryStrings = $queryReflection->resolveQueryStrings($queryExpr, $scope);
         } else {
             $parameterTypes = $scope->getType($paramsExpr);
-            $queryString = $queryReflection->resolvePreparedQueryStrings($queryExpr, $parameterTypes, $scope);
+            $queryStrings = $queryReflection->resolvePreparedQueryStrings($queryExpr, $parameterTypes, $scope);
         }
 
-        return $doctrineReflection->createFetchType($queryString, $methodReflection);
+        return $doctrineReflection->createFetchType($queryStrings, $methodReflection);
     }
 }

src/Extensions/DoctrineResultDynamicReturnTypeExtension.php

@@ -56,33 +56,30 @@ public function getTypeFromMethodCall(MethodReflection $methodReflection, Method
         }
 
         $resultType = $scope->getType($methodCall->var);
-        if (!$resultType instanceof GenericObjectType) {
-            return $defaultReturn;
-        }
-
-        $genericTypes = $resultType->getTypes();
-
-        if (1 !== \count($genericTypes)) {
-            return $defaultReturn;
-        }
-
-        $resultRowType = $genericTypes[0];
-
         if ('columncount' === strtolower($methodReflection->getName())) {
-            if ($resultRowType instanceof ConstantArrayType) {
-                $columnCount = \count($resultRowType->getKeyTypes()) / 2;
-                if (!\is_int($columnCount)) {
-                    throw new ShouldNotHappenException();
+            if ($resultType instanceof GenericObjectType) {
+                $genericTypes = $resultType->getTypes();
+                if (1 !== \count($genericTypes)) {
+                    return $defaultReturn;
                 }
+                $resultRowType = $genericTypes[0];
+
+                if ($resultRowType instanceof ConstantArrayType) {
+                    $columnCount = \count($resultRowType->getKeyTypes()) / 2;
+                    if (!\is_int($columnCount)) {
+                        throw new ShouldNotHappenException();
+                    }
 
-                return new ConstantIntegerType($columnCount);
+                    return new ConstantIntegerType($columnCount);
+                }
             }
 
             return $defaultReturn;
         }
 
         $doctrineReflection = new DoctrineReflection();
-        $fetchResultType = $doctrineReflection->reduceResultType($methodReflection, $resultRowType);
+        $fetchResultType = $doctrineReflection->reduceResultType($methodReflection, $resultType);
+
         if (null !== $fetchResultType) {
             return $fetchResultType;
         }

tests/default/data/doctrine-dbal-union-result.php

@@ -21,9 +21,11 @@ public function doFoo(Connection $conn)
             $result = $stmt->execute([]);
             assertType('Doctrine\DBAL\Result<array{adaid: int<0, 4294967295>, 0: int<0, 4294967295>}>|Doctrine\DBAL\Result<array{email: string, 0: string}>', $result);
 
-            // XXX todo
-            // $fetch = $result->fetchOne();
-            // assertType('Doctrine\DBAL\Result<array{adaid: int<0, 4294967295>, 0: int<0, 4294967295>}>', $fetch);
+            $fetch = $result->fetchOne();
+            assertType('int<0, 4294967295>|string|false', $fetch);
+
+            $fetch = $result->fetchAssociative();
+            assertType('array{adaid: int<0, 4294967295>}|array{email: string}|false', $fetch);
         }
     }