PDO: support PDOStatement::setFetchMode()

config/extensions.neon

@@ -50,6 +50,11 @@ services:
         	- phpstan.broker.dynamicMethodReturnTypeExtension
 
     -
+        class: staabm\PHPStanDba\Extensions\PdoStatementSetFetchModeTypeSpecifyingExtension
+        tags:
+        	- phpstan.typeSpecifier.methodTypeSpecifyingExtension
+
+    -     
         class: staabm\PHPStanDba\Extensions\PdoStatementFetchObjectDynamicReturnTypeExtension
         tags:
         	- phpstan.broker.dynamicMethodReturnTypeExtension

src/Extensions/PdoStatementSetFetchModeTypeSpecifyingExtension.php

@@ -0,0 +1,73 @@
+<?php
+
+declare(strict_types=1);
+
+namespace staabm\PHPStanDba\Extensions;
+
+use PDOStatement;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Analyser\SpecifiedTypes;
+use PHPStan\Analyser\TypeSpecifier;
+use PHPStan\Analyser\TypeSpecifierAwareExtension;
+use PHPStan\Analyser\TypeSpecifierContext;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\MethodTypeSpecifyingExtension;
+use staabm\PHPStanDba\PdoReflection\PdoStatementReflection;
+
+final class PdoStatementSetFetchModeTypeSpecifyingExtension implements MethodTypeSpecifyingExtension, TypeSpecifierAwareExtension
+{
+    private TypeSpecifier $typeSpecifier;
+
+    public function getClass(): string
+    {
+        return PDOStatement::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection, MethodCall $node, TypeSpecifierContext $context): bool
+    {
+        return 'setfetchmode' === strtolower($methodReflection->getName());
+    }
+
+    public function setTypeSpecifier(TypeSpecifier $typeSpecifier): void
+    {
+        $this->typeSpecifier = $typeSpecifier;
+    }
+
+    public function specifyTypes(MethodReflection $methodReflection, MethodCall $node, Scope $scope, TypeSpecifierContext $context): SpecifiedTypes
+    {
+        // keep original param name because named-parameters
+        $methodCall = $node;
+        $statementType = $scope->getType($methodCall->var);
+
+        if ($statementType instanceof GenericObjectType) {
+            $reducedType = $this->reduceType($methodCall, $statementType, $scope);
+
+            if (null !== $reducedType) {
+                return $this->typeSpecifier->create($methodCall->var, $reducedType, TypeSpecifierContext::createTruthy(), true);
+            }
+        }
+
+        return new SpecifiedTypes();
+    }
+
+    private function reduceType(MethodCall $methodCall, GenericObjectType $statementType, Scope $scope): ?GenericObjectType
+    {
+        $args = $methodCall->getArgs();
+
+        if (\count($args) < 1) {
+            return null;
+        }
+
+        $pdoStatementReflection = new PdoStatementReflection();
+
+        $fetchModeType = $scope->getType($args[0]->value);
+        $fetchType = $pdoStatementReflection->getFetchType($fetchModeType);
+        if (null === $fetchType) {
+            return null;
+        }
+
+        return $pdoStatementReflection->modifyGenericStatement($statementType, $fetchType);
+    }
+}

src/PdoReflection/PdoStatementReflection.php

@@ -87,10 +87,10 @@ public function getFetchType(Type $fetchModeType): ?int
      */
     public function createGenericStatement(iterable $queryStrings, int $reflectionFetchType): ?Type
     {
+        $queryReflection = new QueryReflection();
         $genericObjects = [];
 
         foreach ($queryStrings as $queryString) {
-            $queryReflection = new QueryReflection();
             $bothType = $queryReflection->getResultType($queryString, QueryReflector::FETCH_TYPE_BOTH);
 
             if ($bothType) {
@@ -110,6 +110,23 @@ public function createGenericStatement(iterable $queryStrings, int $reflectionFe
         return null;
     }
 
+    /**
+     * @param QueryReflector::FETCH_TYPE* $fetchType
+     */
+    public function modifyGenericStatement(GenericObjectType $statementType, int $fetchType): ?GenericObjectType
+    {
+        $genericTypes = $statementType->getTypes();
+
+        if (2 !== \count($genericTypes)) {
+            return null;
+        }
+
+        $bothType = $genericTypes[1];
+        $rowTypeInFetchMode = $this->reduceStatementResultType($bothType, $fetchType);
+
+        return new GenericObjectType(PDOStatement::class, [$rowTypeInFetchMode, $bothType]);
+    }
+
     /**
      * @param QueryReflector::FETCH_TYPE* $fetchType
      */

tests/default/DbaInferenceTest.php

@@ -2,6 +2,8 @@
 
 namespace staabm\PHPStanDba\Tests;
 
+use Composer\InstalledVersions;
+use Composer\Semver\VersionParser;
 use PHPStan\Testing\TypeInferenceTestCase;
 
 class DbaInferenceTest extends TypeInferenceTestCase
@@ -46,6 +48,11 @@ public function dataFileAsserts(): iterable
         }
 
         yield from $this->gatherAssertTypes(__DIR__.'/data/bug254.php');
+
+        if (InstalledVersions::satisfies(new VersionParser(), 'phpstan/phpstan', '^1.4.7')) {
+            yield from $this->gatherAssertTypes(__DIR__.'/data/pdo-stmt-set-fetch-mode.php');
+        }
+
         yield from $this->gatherAssertTypes(__DIR__.'/data/pdo-union-result.php');
         yield from $this->gatherAssertTypes(__DIR__.'/data/mysqli-union-result.php');
     }

tests/default/data/pdo-stmt-set-fetch-mode.php

@@ -0,0 +1,54 @@
+<?php
+
+namespace PdoStmtFetchModeTest;
+
+use PDO;
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+    public function setFetchModeNum(PDO $pdo)
+    {
+        $bothType = ', array{email: string, 0: string, adaid: int<0, 4294967295>, 1: int<0, 4294967295>}';
+
+        $query = 'SELECT email, adaid FROM ada';
+        $stmt = $pdo->query($query);
+        assertType('PDOStatement<array{email: string, 0: string, adaid: int<0, 4294967295>, 1: int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $stmt->setFetchMode(PDO::FETCH_NUM);
+        assertType('PDOStatement<array{string, int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $result = $stmt->fetch(PDO::FETCH_NUM);
+        assertType('array{string, int<0, 4294967295>}|false', $result);
+    }
+
+    public function setFetchModeAssoc(PDO $pdo)
+    {
+        $bothType = ', array{email: string, 0: string, adaid: int<0, 4294967295>, 1: int<0, 4294967295>}';
+
+        $query = 'SELECT email, adaid FROM ada';
+        $stmt = $pdo->query($query);
+        assertType('PDOStatement<array{email: string, 0: string, adaid: int<0, 4294967295>, 1: int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $stmt->setFetchMode(PDO::FETCH_ASSOC);
+        assertType('PDOStatement<array{email: string, adaid: int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $result = $stmt->fetch(PDO::FETCH_ASSOC);
+        assertType('array{email: string, adaid: int<0, 4294967295>}|false', $result);
+    }
+
+    public function setFetchModeOnQuery(PDO $pdo)
+    {
+        $bothType = ', array{email: string, 0: string, adaid: int<0, 4294967295>, 1: int<0, 4294967295>}';
+
+        $query = 'SELECT email, adaid FROM ada';
+        $stmt = $pdo->query($query, PDO::FETCH_NUM);
+        assertType('PDOStatement<array{string, int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $stmt->setFetchMode(PDO::FETCH_ASSOC);
+        assertType('PDOStatement<array{email: string, adaid: int<0, 4294967295>}'.$bothType.'>', $stmt);
+
+        $result = $stmt->fetch(PDO::FETCH_NUM);
+        assertType('array{string, int<0, 4294967295>}|false', $result);
+    }
+}