[mypyc] feat: implement try_getting_literal in preallocate helper [1/1]

mypyc/irbuild/for_helpers.py

@@ -10,6 +10,7 @@
 from collections.abc import Callable
 from typing import ClassVar, cast
 
+from mypy.checkexpr import try_getting_literal
 from mypy.nodes import (
     ARG_POS,
     CallExpr,
@@ -268,7 +269,8 @@ def sequence_from_generator_preallocate_helper(
                         if isinstance(typ, LiteralType)
                         else TupleGet(sequence, i, line)
                     )
-                    for i, typ in enumerate(proper_types)
+                    # TODO: I think after pulling in recent changes I need to figure out how to handle this change in the `if` block
+                    for i, typ in enumerate(map(try_getting_literal, proper_type.items))
                 ]
 
             items = list(map(builder.add, get_item_ops))

mypyc/test-data/irbuild-tuple.test

@@ -784,71 +784,128 @@ L4:
     a = r8
     return 1
 
-[case testTupleBuiltFromFinalFixedLengthTuple]
+[case testTupleBuiltFromModLevelFixedLengthTuple]
 from typing import Final
 
-source: Final = (1, 2, 3)
+def varint() -> int:
+    # this helper lets us break constant folding for this test
+    return 2
+
+source_final: Final = (1, varint(), 3)
+source_var = (1, varint(), 3)
 
 def f(val: int) -> bool:
     return val % 2 == 0
 
-def test() -> None:
-    a = tuple(f(x) for x in source)
+def test_final_can_fold() -> None:
+    a = tuple(f(x) for x in source_final)
+def test_var_can_not_fold() -> None:
+    a = tuple(f(x) for x in source_var)
 [out]
+def varint():
+L0:
+    return 4
 def f(val):
     val, r0 :: int
     r1 :: bit
 L0:
     r0 = CPyTagged_Remainder(val, 4)
     r1 = int_eq r0, 0
     return r1
-def test():
+def test_final_can_fold():
     r0 :: tuple[int, int, int]
     r1 :: bool
-    r2, r3, r4 :: int
-    r5, r6, r7 :: object
-    r8, r9 :: tuple
-    r10 :: native_int
-    r11 :: bit
-    r12 :: object
-    r13, x :: int
-    r14 :: bool
-    r15 :: object
-    r16 :: native_int
+    r2 :: object
+    r3 :: int
+    r4, r5 :: object
+    r6, r7 :: tuple
+    r8 :: native_int
+    r9 :: bit
+    r10 :: object
+    r11, x :: int
+    r12 :: bool
+    r13 :: object
+    r14 :: native_int
     a :: tuple
 L0:
-    r0 = __main__.source :: static
+    r0 = __main__.source_final :: static
     if is_error(r0) goto L1 else goto L2
 L1:
-    r1 = raise NameError('value for final name "source" was not set')
+    r1 = raise NameError('value for final name "source_final" was not set')
     unreachable
 L2:
-    r2 = r0[0]
+    r2 = object 1
     r3 = r0[1]
-    r4 = r0[2]
-    r5 = box(int, r2)
-    r6 = box(int, r3)
-    r7 = box(int, r4)
-    r8 = PyTuple_Pack(3, r5, r6, r7)
-    r9 = PyTuple_New(3)
-    r10 = 0
+    r4 = object 3
+    r5 = box(int, r3)
+    r6 = PyTuple_Pack(3, r2, r5, r4)
+    r7 = PyTuple_New(3)
+    r8 = 0
     goto L4
 L3:
-    r11 = r10 < 3 :: signed
-    if r11 goto L4 else goto L6 :: bool
+    r9 = r8 < 3 :: signed
+    if r9 goto L4 else goto L6 :: bool
 L4:
-    r12 = CPySequenceTuple_GetItemUnsafe(r8, r10)
-    r13 = unbox(int, r12)
-    x = r13
-    r14 = f(x)
-    r15 = box(bool, r14)
-    CPySequenceTuple_SetItemUnsafe(r9, r10, r15)
+    r10 = CPySequenceTuple_GetItemUnsafe(r6, r8)
+    r11 = unbox(int, r10)
+    x = r11
+    r12 = f(x)
+    r13 = box(bool, r12)
+    CPySequenceTuple_SetItemUnsafe(r7, r8, r13)
 L5:
-    r16 = r10 + 1
-    r10 = r16
+    r14 = r8 + 1
+    r8 = r14
     goto L3
 L6:
-    a = r9
+    a = r7
+    return 1
+def test_var_can_not_fold():
+    r0 :: dict
+    r1 :: str
+    r2 :: object
+    r3 :: tuple[int, int, int]
+    r4, r5, r6 :: int
+    r7, r8, r9 :: object
+    r10, r11 :: tuple
+    r12 :: native_int
+    r13 :: bit
+    r14 :: object
+    r15, x :: int
+    r16 :: bool
+    r17 :: object
+    r18 :: native_int
+    a :: tuple
+L0:
+    r0 = __main__.globals :: static
+    r1 = 'source_var'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = unbox(tuple[int, int, int], r2)
+    r4 = r3[0]
+    r5 = r3[1]
+    r6 = r3[2]
+    r7 = box(int, r4)
+    r8 = box(int, r5)
+    r9 = box(int, r6)
+    r10 = PyTuple_Pack(3, r7, r8, r9)
+    r11 = PyTuple_New(3)
+    r12 = 0
+    goto L2
+L1:
+    r13 = r12 < 3 :: signed
+    if r13 goto L2 else goto L4 :: bool
+L2:
+    r14 = CPySequenceTuple_GetItemUnsafe(r10, r12)
+    r15 = unbox(int, r14)
+    x = r15
+    r16 = f(x)
+    r17 = box(bool, r16)
+    CPySequenceTuple_SetItemUnsafe(r11, r12, r17)
+L3:
+    r18 = r12 + 1
+    r12 = r18
+    goto L1
+L4:
+    a = r11
     return 1
 
 [case testTupleBuiltFromVariableLengthTuple]