Create keyed arrays when assigning literal union keys

Author: danog

src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php

@@ -349,29 +349,23 @@ private static function updateTypeWithKeyValues(
         }
 
         if (!$has_matching_objectlike_property && !$has_matching_string) {
-            if (count($key_values) === 1) {
-                $key_value = $key_values[0];
-
-                $object_like = new TKeyedArray(
-                    [$key_value->value => $current_type],
-                    $key_value instanceof TLiteralClassString
-                        ? [$key_value->value => true]
-                        : null,
-                );
-
-                $array_assignment_type = new Union([
-                    $object_like,
-                ]);
-            } else {
-                $array_assignment_literals = $key_values;
-
-                $array_assignment_type = new Union([
-                    new TNonEmptyArray([
-                        new Union($array_assignment_literals),
-                        $current_type,
-                    ]),
-                ]);
+            $properties = [];
+            $classStrings = [];
+            $current_type = $current_type->setPossiblyUndefined(count($key_values) > 1);
+            foreach ($key_values as $key_value) {
+                $properties[$key_value->value] = $current_type;
+                if ($key_value instanceof TLiteralClassString) {
+                    $classStrings[$key_value->value] = true;
+                }
             }
+            $object_like = new TKeyedArray(
+                $properties,
+                $classStrings ?: null,
+            );
+
+            $array_assignment_type = new Union([
+                $object_like,
+            ]);
 
             return Type::combineUnionTypes(
                 $child_stmt_type,

tests/ArrayAssignmentTest.php

@@ -34,6 +34,27 @@ public function testConditionalAssignment(): void
     public function providerValidCodeParse(): iterable
     {
         return [
+            'assignUnionOfLiterals' => [
+                'code' => '<?php
+                    $result = [];
+                    
+                    foreach (["a", "b"] as $k) {
+                        $result[$k] = true;
+                    }
+                    
+                    $resultOpt = [];
+                    
+                    foreach (["a", "b"] as $k) {
+                        if (random_int(0, 1)) {
+                            continue;
+                        }
+                        $resultOpt[$k] = true;
+                    }',
+                'assertions' => [
+                    '$result===' => 'array{a: true, b: true}',
+                    '$resultOpt===' => 'array{a?: true, b?: true}',
+                ],
+            ],
             'genericArrayCreationWithSingleIntValue' => [
                 'code' => '<?php
                     $out = [];