From bb1a89ecb76522eca16facece541c3d9451420c0 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Wed, 16 Oct 2024 12:17:35 +0300
Subject: [PATCH] [python] When inferring TypedDict from dict literals,
 preserve the order of their value types

For instance, for `{"foo": 42, "bar": "baz"}` infer `dict[str, int | str]`, not `dict[str, str | int]`.

GitOrigin-RevId: c0063639c2122054904be7dd242d9b1a6fc2fbc8
---
 .../python/psi/types/PyCollectionTypeUtil.kt  |  2 +-
 .../annotationInferredTypedDictPre39_after.py |  2 +-
 .../annotationInferredTypedDict_after.py      |  2 +-
 .../PyTypeCheckerInspectionTest.java          | 29 ++++++++++---------
 4 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCollectionTypeUtil.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCollectionTypeUtil.kt
index b4c2ec261cbd..45efca7c70f3 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCollectionTypeUtil.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCollectionTypeUtil.kt
@@ -106,7 +106,7 @@ object PyCollectionTypeUtil {
     val elements = sequence.elements
     val maxAnalyzedElements = MAX_ANALYZED_ELEMENTS_OF_LITERALS.coerceAtMost(elements.size)
     var allStrKeys = true
-    val strKeysToValueTypes = HashMap<String, Pair<PyExpression?, PyType?>>()
+    val strKeysToValueTypes = LinkedHashMap<String, Pair<PyExpression?, PyType?>>()
 
     elements
       .take(maxAnalyzedElements)
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDictPre39_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDictPre39_after.py
index 06c99fda638f..ff63078bda63 100644
--- a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDictPre39_after.py
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDictPre39_after.py
@@ -1,3 +1,3 @@
 from typing import Dict, Union
 
-var: [Dict[str, Union[str, int]]] = {"foo": 42, "bar": "baz"}
\ No newline at end of file
+var: [Dict[str, Union[int, str]]] = {"foo": 42, "bar": "baz"}
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDict_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDict_after.py
index b1896baaf7be..438d18fc212a 100644
--- a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDict_after.py
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationInferredTypedDict_after.py
@@ -1 +1 @@
-var: [dict[str, str | int]] = {"foo": 42, "bar": "baz"}
\ No newline at end of file
+var: [dict[str, int | str]] = {"foo": 42, "bar": "baz"}
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
index e2142abc072c..76dc6b70705b 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
@@ -1414,14 +1414,15 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
                            z = {'foo': 'bar'}
                            n = {"foo": "", "quux": 3}
                            f(<warning descr="Expected type 'C', got 'dict' instead">y</warning>)
-                           f(<warning descr="Expected type 'C', got 'dict[str, int | str]' instead">n</warning>)
+                           f(<warning descr="Expected type 'C', got 'dict[str, str | int]' instead">n</warning>)
                            f(z)
                            f(x=<warning descr="Expected type 'C', got 'dict' instead">y</warning>)
-                           f(x=<warning descr="Expected type 'C', got 'dict[str, int | str]' instead">n</warning>)
+                           f(x=<warning descr="Expected type 'C', got 'dict[str, str | int]' instead">n</warning>)
                            f(x=z)
                            z2: C = <warning descr="Expected type 'C', got 'dict' instead">y</warning>
-                           z2: C = <warning descr="Expected type 'C', got 'dict[str, int | str]' instead">n</warning>
-                           z2: C = z""")
+                           z2: C = <warning descr="Expected type 'C', got 'dict[str, str | int]' instead">n</warning>
+                           z2: C = z
+                           """)
     );
   }
 
@@ -1447,18 +1448,18 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
                                a: str
                                b: str
                                c: str
-
-
+                           
+                           
                            class NotSoHardDict(TypedDict):
                                a: str
                                b: EasyDict
-
-
+                           
+                           
                            class HardDict(TypedDict):
                                a: str
                                d: NotSoHardDict
-
-
+                           
+                           
                            q: HardDict = {
                                'a': <warning descr="Expected type 'str', got 'int' instead">42</warning>,
                                'd': {
@@ -1473,8 +1474,8 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
                                    0: 'zero',
                                }
                            }
-                           s: HardDict = {'a': 'xx', 'd': <warning descr="Expected type 'NotSoHardDict', got 'dict[str, dict[int, str] | str]' instead">t</warning>}
-                           s1: HardDict = <warning descr="Expected type 'HardDict', got 'dict[str, dict[int, str] | str]' instead">t</warning>
+                           s: HardDict = {'a': 'xx', 'd': <warning descr="Expected type 'NotSoHardDict', got 'dict[str, str | dict[int, str]]' instead">t</warning>}
+                           s1: HardDict = <warning descr="Expected type 'HardDict', got 'dict[str, str | dict[int, str]]' instead">t</warning>
                            t1 = {
                                'a': 'xx',
                                'd': {
@@ -1482,8 +1483,8 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
                                    'd': {}
                                }
                            }
-                           s2: HardDict = {'a': 'xx', 'd': <warning descr="Expected type 'NotSoHardDict', got 'dict[str, dict[str, dict | int] | str]' instead">t1</warning>}
-                           s3: HardDict = <warning descr="Expected type 'HardDict', got 'dict[str, dict[str, dict | int] | str]' instead">t1</warning>
+                           s2: HardDict = {'a': 'xx', 'd': <warning descr="Expected type 'NotSoHardDict', got 'dict[str, str | dict[str, int | dict]]' instead">t1</warning>}
+                           s3: HardDict = <warning descr="Expected type 'HardDict', got 'dict[str, str | dict[str, int | dict]]' instead">t1</warning>
                            s4: HardDict = <warning descr="TypedDict 'HardDict' has missing key: 'a'">{
                                'd': {
                                    'a': 'a',