2929 TypeAlias ,
3030 Var ,
3131)
32- from mypy .types import Type
32+ from mypy .types import LiteralType , TupleType , Type , get_proper_type , get_proper_types
3333from mypyc .ir .ops import (
3434 ERR_NEVER ,
3535 BasicBlock ,
3838 IntOp ,
3939 LoadAddress ,
4040 LoadErrorValue ,
41+ LoadLiteral ,
4142 LoadMem ,
4243 MethodCall ,
4344 RaiseStandardError ,
@@ -172,7 +173,7 @@ def for_loop_helper_with_index(
172173 body_insts: a function that generates the body of the loop.
173174 It needs a index as parameter.
174175 """
175- assert is_sequence_rprimitive (expr_reg .type )
176+ assert is_sequence_rprimitive (expr_reg .type ), ( expr_reg , expr_reg . type )
176177 target_type = builder .get_sequence_type (expr )
177178
178179 body_block = BasicBlock ()
@@ -219,10 +220,9 @@ def sequence_from_generator_preallocate_helper(
219220 there is no condition list in the generator and only one original sequence with
220221 one index is allowed.
221222
222- e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes)
223- (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes)
224- (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes]
225- RTuple as an original sequence is not supported yet.
223+ e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
224+ (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
225+ (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple]
226226
227227 Args:
228228 empty_op_llbuilder: A function that can generate an empty sequence op when
@@ -237,23 +237,41 @@ def sequence_from_generator_preallocate_helper(
237237 implementation.
238238 """
239239 if len (gen .sequences ) == 1 and len (gen .indices ) == 1 and len (gen .condlists [0 ]) == 0 :
240- rtype = builder .node_type (gen .sequences [0 ])
241- if is_sequence_rprimitive (rtype ):
242- sequence = builder .accept (gen .sequences [0 ])
243- length = get_expr_length_value (
244- builder , gen .sequences [0 ], sequence , gen .line , use_pyssize_t = True
245- )
246- target_op = empty_op_llbuilder (length , gen .line )
247-
248- def set_item (item_index : Value ) -> None :
249- e = builder .accept (gen .left_expr )
250- builder .call_c (set_item_op , [target_op , item_index , e ], gen .line )
251-
252- for_loop_helper_with_index (
253- builder , gen .indices [0 ], gen .sequences [0 ], sequence , set_item , gen .line , length
254- )
240+ line = gen .line
241+ sequence_expr = gen .sequences [0 ]
242+ rtype = builder .node_type (sequence_expr )
243+ if not (is_sequence_rprimitive (rtype ) or isinstance (rtype , RTuple )):
244+ return None
245+ sequence = builder .accept (sequence_expr )
246+ length = get_expr_length_value (builder , sequence_expr , sequence , line , use_pyssize_t = True )
247+ if isinstance (rtype , RTuple ):
248+ # If input is RTuple, box it to tuple_rprimitive for generic iteration
249+ # TODO: this can be optimized a bit better with an unrolled ForRTuple helper
250+ proper_type = get_proper_type (builder .types [sequence_expr ])
251+ assert isinstance (proper_type , TupleType ), proper_type
252+
253+ get_item_ops = [
254+ (
255+ LoadLiteral (typ .value , object_rprimitive )
256+ if isinstance (typ , LiteralType )
257+ else TupleGet (sequence , i , line )
258+ )
259+ for i , typ in enumerate (get_proper_types (proper_type .items ))
260+ ]
261+ items = list (map (builder .add , get_item_ops ))
262+ sequence = builder .new_tuple (items , line )
263+
264+ target_op = empty_op_llbuilder (length , line )
265+
266+ def set_item (item_index : Value ) -> None :
267+ e = builder .accept (gen .left_expr )
268+ builder .call_c (set_item_op , [target_op , item_index , e ], line )
269+
270+ for_loop_helper_with_index (
271+ builder , gen .indices [0 ], sequence_expr , sequence , set_item , line , length
272+ )
255273
256- return target_op
274+ return target_op
257275 return None
258276
259277
@@ -806,7 +824,7 @@ class ForSequence(ForGenerator):
806824 def init (
807825 self , expr_reg : Value , target_type : RType , reverse : bool , length : Value | None = None
808826 ) -> None :
809- assert is_sequence_rprimitive (expr_reg .type ), expr_reg
827+ assert is_sequence_rprimitive (expr_reg .type ), ( expr_reg , expr_reg . type )
810828 builder = self .builder
811829 # Record a Value indicating the length of the sequence, if known at compile time.
812830 self .length = length
@@ -836,7 +854,6 @@ def init(
836854 def gen_condition (self ) -> None :
837855 builder = self .builder
838856 line = self .line
839- # TODO: Don't reload the length each time when iterating an immutable sequence?
840857 if self .reverse :
841858 # If we are iterating in reverse order, we obviously need
842859 # to check that the index is still positive. Somewhat less
@@ -1218,7 +1235,7 @@ def get_expr_length_value(
12181235 builder : IRBuilder , expr : Expression , expr_reg : Value , line : int , use_pyssize_t : bool
12191236) -> Value :
12201237 rtype = builder .node_type (expr )
1221- assert is_sequence_rprimitive (rtype ), rtype
1238+ assert is_sequence_rprimitive (rtype ) or isinstance ( rtype , RTuple ) , rtype
12221239 length = get_expr_length (expr )
12231240 if length is None :
12241241 # We cannot compute the length at compile time, so we will fetch it.
0 commit comments