2525 TupleExpr ,
2626 TypeAlias ,
2727)
28+ from mypy .types import LiteralType , TupleType , get_proper_type , get_proper_types
2829from mypyc .ir .ops import (
2930 ERR_NEVER ,
3031 BasicBlock ,
3334 IntOp ,
3435 LoadAddress ,
3536 LoadErrorValue ,
37+ LoadLiteral ,
3638 LoadMem ,
3739 MethodCall ,
3840 RaiseStandardError ,
@@ -168,7 +170,7 @@ def for_loop_helper_with_index(
168170 body_insts: a function that generates the body of the loop.
169171 It needs a index as parameter.
170172 """
171- assert is_sequence_rprimitive (expr_reg .type )
173+ assert is_sequence_rprimitive (expr_reg .type ), ( expr_reg , expr_reg . type )
172174 target_type = builder .get_sequence_type (expr )
173175
174176 body_block = BasicBlock ()
@@ -215,10 +217,9 @@ def sequence_from_generator_preallocate_helper(
215217 there is no condition list in the generator and only one original sequence with
216218 one index is allowed.
217219
218- e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes)
219- (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes)
220- (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes]
221- RTuple as an original sequence is not supported yet.
220+ e.g. (1) tuple(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
221+ (2) list(f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple)
222+ (3) [f(x) for x in a_list/a_tuple/a_str/a_bytes/an_rtuple]
222223
223224 Args:
224225 empty_op_llbuilder: A function that can generate an empty sequence op when
@@ -233,23 +234,41 @@ def sequence_from_generator_preallocate_helper(
233234 implementation.
234235 """
235236 if len (gen .sequences ) == 1 and len (gen .indices ) == 1 and len (gen .condlists [0 ]) == 0 :
236- rtype = builder .node_type (gen .sequences [0 ])
237- if is_sequence_rprimitive (rtype ):
238- sequence = builder .accept (gen .sequences [0 ])
239- length = get_expr_length_value (
240- builder , gen .sequences [0 ], sequence , gen .line , use_pyssize_t = True
241- )
242- target_op = empty_op_llbuilder (length , gen .line )
243-
244- def set_item (item_index : Value ) -> None :
245- e = builder .accept (gen .left_expr )
246- builder .call_c (set_item_op , [target_op , item_index , e ], gen .line )
247-
248- for_loop_helper_with_index (
249- builder , gen .indices [0 ], gen .sequences [0 ], sequence , set_item , gen .line , length
250- )
237+ line = gen .line
238+ sequence_expr = gen .sequences [0 ]
239+ rtype = builder .node_type (sequence_expr )
240+ if not (is_sequence_rprimitive (rtype ) or isinstance (rtype , RTuple )):
241+ return None
242+ sequence = builder .accept (sequence_expr )
243+ length = get_expr_length_value (builder , sequence_expr , sequence , line , use_pyssize_t = True )
244+ if isinstance (rtype , RTuple ):
245+ # If input is RTuple, box it to tuple_rprimitive for generic iteration
246+ # TODO: this can be optimized a bit better with an unrolled ForRTuple helper
247+ proper_type = get_proper_type (builder .types [sequence_expr ])
248+ assert isinstance (proper_type , TupleType ), proper_type
249+
250+ get_item_ops = [
251+ (
252+ LoadLiteral (typ .value , object_rprimitive )
253+ if isinstance (typ , LiteralType )
254+ else TupleGet (sequence , i , line )
255+ )
256+ for i , typ in enumerate (get_proper_types (proper_type .items ))
257+ ]
258+ items = list (map (builder .add , get_item_ops ))
259+ sequence = builder .new_tuple (items , line )
260+
261+ target_op = empty_op_llbuilder (length , line )
262+
263+ def set_item (item_index : Value ) -> None :
264+ e = builder .accept (gen .left_expr )
265+ builder .call_c (set_item_op , [target_op , item_index , e ], line )
266+
267+ for_loop_helper_with_index (
268+ builder , gen .indices [0 ], sequence_expr , sequence , set_item , line , length
269+ )
251270
252- return target_op
271+ return target_op
253272 return None
254273
255274
@@ -802,7 +821,7 @@ class ForSequence(ForGenerator):
802821 def init (
803822 self , expr_reg : Value , target_type : RType , reverse : bool , length : Value | None = None
804823 ) -> None :
805- assert is_sequence_rprimitive (expr_reg .type ), expr_reg
824+ assert is_sequence_rprimitive (expr_reg .type ), ( expr_reg , expr_reg . type )
806825 builder = self .builder
807826 # Record a Value indicating the length of the sequence, if known at compile time.
808827 self .length = length
@@ -832,7 +851,6 @@ def init(
832851 def gen_condition (self ) -> None :
833852 builder = self .builder
834853 line = self .line
835- # TODO: Don't reload the length each time when iterating an immutable sequence?
836854 if self .reverse :
837855 # If we are iterating in reverse order, we obviously need
838856 # to check that the index is still positive. Somewhat less
@@ -1207,7 +1225,7 @@ def get_expr_length_value(
12071225 builder : IRBuilder , expr : Expression , expr_reg : Value , line : int , use_pyssize_t : bool
12081226) -> Value :
12091227 rtype = builder .node_type (expr )
1210- assert is_sequence_rprimitive (rtype ), rtype
1228+ assert is_sequence_rprimitive (rtype ) or isinstance ( rtype , RTuple ) , rtype
12111229 length = get_expr_length (builder , expr )
12121230 if length is None :
12131231 # We cannot compute the length at compile time, so we will fetch it.
0 commit comments