Merge pull request #7449 from orklah/7415-2

Array key exists assert both ways
vimeo · Jan 20, 2022 · f6369dc · f6369dc
2 parents 831f14d + 35577df
commit f6369dc
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 49 deletions.
diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/AssertionFinder.php b/src/Psalm/Internal/Analyzer/Statements/Expression/AssertionFinder.php
@@ -3698,66 +3698,66 @@ private static function getArrayKeyExistsAssertions(
 
         if ($literal_assertions && $first_var_name) {
             $if_types[$first_var_name] = [$literal_assertions];
-        } else {
-            $array_root = isset($expr->getArgs()[1]->value)
-                ? ExpressionIdentifier::getArrayVarId(
-                    $expr->getArgs()[1]->value,
-                    $this_class_name,
-                    $source
-                )
-                : null;
+        }
+
+        $array_root = isset($expr->getArgs()[1]->value)
+            ? ExpressionIdentifier::getArrayVarId(
+                $expr->getArgs()[1]->value,
+                $this_class_name,
+                $source
+            )
+            : null;
 
-            if ($array_root) {
-                if ($first_var_name === null && isset($expr->getArgs()[0])) {
-                    $first_arg = $expr->getArgs()[0];
+        if ($array_root) {
+            if ($first_var_name === null && isset($expr->getArgs()[0])) {
+                $first_arg = $expr->getArgs()[0];
 
-                    if ($first_arg->value instanceof PhpParser\Node\Scalar\String_) {
-                        $first_var_name = '\'' . $first_arg->value->value . '\'';
-                    } elseif ($first_arg->value instanceof PhpParser\Node\Scalar\LNumber) {
-                        $first_var_name = (string)$first_arg->value->value;
-                    }
+                if ($first_arg->value instanceof PhpParser\Node\Scalar\String_) {
+                    $first_var_name = '\'' . $first_arg->value->value . '\'';
+                } elseif ($first_arg->value instanceof PhpParser\Node\Scalar\LNumber) {
+                    $first_var_name = (string)$first_arg->value->value;
                 }
+            }
 
-                if ($expr->getArgs()[0]->value instanceof PhpParser\Node\Expr\ClassConstFetch
-                    && $expr->getArgs()[0]->value->name instanceof PhpParser\Node\Identifier
-                    && $expr->getArgs()[0]->value->name->name !== 'class'
-                ) {
-                    $const_type = null;
+            if ($expr->getArgs()[0]->value instanceof PhpParser\Node\Expr\ClassConstFetch
+                && $expr->getArgs()[0]->value->name instanceof PhpParser\Node\Identifier
+                && $expr->getArgs()[0]->value->name->name !== 'class'
+            ) {
+                $const_type = null;
 
-                    if ($source instanceof StatementsAnalyzer) {
-                        $const_type = $source->node_data->getType($expr->getArgs()[0]->value);
-                    }
+                if ($source instanceof StatementsAnalyzer) {
+                    $const_type = $source->node_data->getType($expr->getArgs()[0]->value);
+                }
 
-                    if ($const_type) {
-                        if ($const_type->isSingleStringLiteral()) {
-                            $first_var_name = $const_type->getSingleStringLiteral()->value;
-                        } elseif ($const_type->isSingleIntLiteral()) {
-                            $first_var_name = (string)$const_type->getSingleIntLiteral()->value;
-                        } else {
-                            $first_var_name = null;
-                        }
+                if ($const_type) {
+                    if ($const_type->isSingleStringLiteral()) {
+                        $first_var_name = $const_type->getSingleStringLiteral()->value;
+                    } elseif ($const_type->isSingleIntLiteral()) {
+                        $first_var_name = (string)$const_type->getSingleIntLiteral()->value;
                     } else {
                         $first_var_name = null;
                     }
-                } elseif ($expr->getArgs()[0]->value instanceof PhpParser\Node\Expr\Variable
-                    && $source instanceof StatementsAnalyzer
-                    && ($first_var_type = $source->node_data->getType($expr->getArgs()[0]->value))
-                ) {
-                    foreach ($first_var_type->getLiteralStrings() as $array_literal_type) {
-                        $if_types[$array_root . "['" . $array_literal_type->value . "']"] = [[new ArrayKeyExists()]];
-                    }
-                    foreach ($first_var_type->getLiteralInts() as $array_literal_type) {
-                        $if_types[$array_root . "[" . $array_literal_type->value . "]"] = [[new ArrayKeyExists()]];
-                    }
+                } else {
+                    $first_var_name = null;
                 }
-
-                if ($first_var_name !== null
-                    && !strpos($first_var_name, '->')
-                    && !strpos($first_var_name, '[')
-                ) {
-                    $if_types[$array_root . '[' . $first_var_name . ']'] = [[new ArrayKeyExists()]];
+            } elseif ($expr->getArgs()[0]->value instanceof PhpParser\Node\Expr\Variable
+                && $source instanceof StatementsAnalyzer
+                && ($first_var_type = $source->node_data->getType($expr->getArgs()[0]->value))
+            ) {
+                foreach ($first_var_type->getLiteralStrings() as $array_literal_type) {
+                    $if_types[$array_root . "['" . $array_literal_type->value . "']"] = [[new ArrayKeyExists()]];
+                }
+                foreach ($first_var_type->getLiteralInts() as $array_literal_type) {
+                    $if_types[$array_root . "[" . $array_literal_type->value . "]"] = [[new ArrayKeyExists()]];
                 }
             }
+
+            if ($first_var_name !== null
+                && !strpos($first_var_name, '->')
+                && !strpos($first_var_name, '[')
+            ) {
+                $if_types[$array_root . '[' . $first_var_name . ']'] = [[new ArrayKeyExists()]];
+            }
         }
 
         return $if_types ? [$if_types] : [];

diff --git a/src/Psalm/Internal/Codebase/Scanner.php b/src/Psalm/Internal/Codebase/Scanner.php
@@ -80,7 +80,7 @@
 /**
  * @internal
  *
- * Contains methods that aid in the scanning of Psalm's codebase>
+ * Contains methods that aid in the scanning of Psalm's codebase
  */
 class Scanner
 {

diff --git a/tests/TypeReconciliation/ArrayKeyExistsTest.php b/tests/TypeReconciliation/ArrayKeyExistsTest.php
@@ -378,6 +378,35 @@ function checkArrayKeyExistsComparisonNegated(array $arr, string $key): bool
                         return true;
                     }'
             ],
+            'arrayKeyExistsAssertBothWays' => [
+                'code' => '<?php
+                    class a {
+                        const STATE_A = 0;
+                        const STATE_B = 1;
+                        const STATE_C = 2;
+                        /**
+                         * @return array<self::STATE_*, non-empty-string>
+                         * @psalm-pure
+                         */
+                        public static function getStateLabels(): array {
+                            return [
+                                self::STATE_A => "A",
+                                self::STATE_B => "B",
+                                self::STATE_C => "C",
+                            ];
+                        }
+                        /**
+                         * @param self::STATE_* $state
+                         */
+                        public static function getStateLabelIf(int $state): string {
+                            $states = self::getStateLabels();
+                            if (array_key_exists($state, $states)) {
+                                return $states[$state];
+                            }
+                            return "";
+                        }
+                    }'
+            ],
         ];
     }