Add support for most PDOStatement fetch modes (#2553)

Flags like PDO::FETCH_CLASSTYPE and PDO::FETCH_PROPS_LATE aren't supported yet. (in re. #2529)
vimeo · Jan 6, 2020 · 63dea52 · 63dea52
1 parent 3a18b16
commit 63dea52
Show file tree

Hide file tree

Showing 2 changed files with 126 additions and 7 deletions.
diff --git a/src/Psalm/Internal/Provider/ReturnTypeProvider/PdoStatementReturnTypeProvider.php b/src/Psalm/Internal/Provider/ReturnTypeProvider/PdoStatementReturnTypeProvider.php
@@ -36,10 +36,56 @@ public static function getMethodReturnType(
             && ($first_arg_type = $source->getNodeTypeProvider()->getType($call_args[0]->value))
             && $first_arg_type->isSingleIntLiteral()
         ) {
-            $value = $first_arg_type->getSingleIntLiteral()->value;
+            $fetch_mode = $first_arg_type->getSingleIntLiteral()->value;
 
-            if ($value === \PDO::FETCH_ASSOC) {
-                return Type::getArray();
+            switch ($fetch_mode) {
+                case \PDO::FETCH_ASSOC: // array<string,scalar>
+                    return new Type\Union([
+                        new Type\Atomic\TArray([
+                            Type::getString(),
+                            Type::getScalar()
+                        ]),
+                    ]);
+
+                case \PDO::FETCH_BOTH: // array<array-key,scalar>
+                    return new Type\Union([
+                        new Type\Atomic\TArray([
+                            Type::getArrayKey(),
+                            Type::getScalar()
+                        ]),
+                    ]);
+
+                case \PDO::FETCH_BOUND: // true
+                    return Type::getTrue();
+
+                case \PDO::FETCH_CLASS: // object
+                    return Type::getObject();
+
+                case \PDO::FETCH_LAZY: // object
+                    // This actually returns a PDORow object, but that class is
+                    // undocumented, and its attributes are all dynamic anyway
+                    return Type::getObject();
+
+                case \PDO::FETCH_NAMED: // array<string, scalar|list<scalar>>
+                    return new Type\Union([
+                        new Type\Atomic\TArray([
+                            Type::getString(),
+                            new Type\Union([
+                                new Type\Atomic\TScalar(),
+                                new Type\Atomic\TList(Type::getScalar())
+                            ])
+                        ]),
+                    ]);
+
+                case \PDO::FETCH_NUM: // list<scalar>
+                    return new Type\Union([
+                        new Type\Atomic\TList(Type::getScalar())
+                    ]);
+
+                case \PDO::FETCH_OBJ: // stdClass
+                    return new Type\Union([
+                        new Type\Atomic\TNamedObject('stdClass')
+                    ]);
             }
         }
     }

diff --git a/tests/MethodCallTest.php b/tests/MethodCallTest.php
@@ -428,10 +428,83 @@ function foo(string $foo): string {
             ],
             'pdoStatementFetchAssoc' => [
                 '<?php
-                    $p = new \PDO("sqlite:::memory:");
-                    $sth = $p->prepare("SELECT 1");
-                    $sth->execute();
-                    while (($row = $sth->fetch(\PDO::FETCH_ASSOC)) !== false) {}'
+                    /** @return array<string,scalar> */
+                    function fetch_assoc() : array {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_ASSOC);
+                    }'
+            ],
+            'pdoStatementFetchBoth' => [
+                '<?php
+                    /** @return array<scalar> */
+                    function fetch_both() : array {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_BOTH);
+                    }'
+            ],
+            'pdoStatementFetchBound' => [
+                '<?php
+                    /** @return true */
+                    function fetch_both() : bool {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_BOUND);
+                    }'
+            ],
+            'pdoStatementFetchClass' => [
+                '<?php
+                    /** @return object */
+                    function fetch_class() : object {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_CLASS);
+                    }'
+            ],
+            'pdoStatementFetchLazy' => [
+                '<?php
+                    /** @return object */
+                    function fetch_lazy() : object {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_LAZY);
+                    }'
+            ],
+            'pdoStatementFetchNamed' => [
+                '<?php
+                    /** @return array<string,scalar|list<scalar>> */
+                    function fetch_named() : array {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_NAMED);
+                    }'
+            ],
+            'pdoStatementFetchNum' => [
+                '<?php
+                    /** @return list<scalar> */
+                    function fetch_named() : array {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_NUM);
+                    }'
+            ],
+            'pdoStatementFetchObj' => [
+                '<?php
+                    /** @return stdClass */
+                    function fetch_named() : object {
+                        $p = new PDO("sqlite::memory:");
+                        $sth = $p->prepare("SELECT 1");
+                        $sth->execute();
+                        return $sth->fetch(PDO::FETCH_OBJ);
+                    }'
             ],
         ];
     }