Skip to content

Commit 7780482

Browse files
[mypyc] Modify specializers and add int.__pow__ in test-data/fixture (#10448)
* Move simple tuple and list creation to new functions. Each specializer only focus on one situation now. * Add int.__pow__ in test-data/fixture
1 parent 799e40c commit 7780482

File tree

3 files changed

+36
-21
lines changed

3 files changed

+36
-21
lines changed

mypyc/irbuild/specialize.py

Lines changed: 32 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -97,19 +97,6 @@ def dict_methods_fast_path(
9797
if not (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]):
9898
return None
9999
arg = expr.args[0]
100-
101-
# Special case for simplest list comprehension, for example
102-
# list(x for x in tmp_list)
103-
# TODO: The following code should be moved to a new function after
104-
# supporting multiple specialize functions
105-
if not isinstance(callee, MemberExpr) and isinstance(arg, GeneratorExpr):
106-
val = sequence_from_generator_preallocate_helper(
107-
builder, arg,
108-
empty_op_llbuilder=builder.builder.new_list_op_with_length,
109-
set_item_op=new_list_set_item_op)
110-
if val is not None:
111-
return val
112-
113100
if not (isinstance(arg, CallExpr) and not arg.args
114101
and isinstance(arg.callee, MemberExpr)):
115102
return None
@@ -130,6 +117,38 @@ def dict_methods_fast_path(
130117
return builder.call_c(dict_items_op, [obj], expr.line)
131118

132119

120+
@specialize_function('builtins.list')
121+
def translate_list_from_generator_call(
122+
builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
123+
# Special case for simplest list comprehension, for example
124+
# list(f(x) for x in other_list/other_tuple)
125+
# translate_list_comprehension() would take care of other cases if this fails.
126+
if (len(expr.args) == 1
127+
and expr.arg_kinds[0] == ARG_POS
128+
and isinstance(expr.args[0], GeneratorExpr)):
129+
return sequence_from_generator_preallocate_helper(
130+
builder, expr.args[0],
131+
empty_op_llbuilder=builder.builder.new_list_op_with_length,
132+
set_item_op=new_list_set_item_op)
133+
return None
134+
135+
136+
@specialize_function('builtins.tuple')
137+
def translate_tuple_from_generator_call(
138+
builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
139+
# Special case for simplest tuple creation from a generator, for example
140+
# tuple(f(x) for x in other_list/other_tuple)
141+
# translate_safe_generator_call() would take care of other cases if this fails.
142+
if (len(expr.args) == 1
143+
and expr.arg_kinds[0] == ARG_POS
144+
and isinstance(expr.args[0], GeneratorExpr)):
145+
return sequence_from_generator_preallocate_helper(
146+
builder, expr.args[0],
147+
empty_op_llbuilder=builder.builder.new_tuple_with_length,
148+
set_item_op=new_tuple_set_item_op)
149+
return None
150+
151+
133152
@specialize_function('builtins.set')
134153
def translate_set_from_generator_call(
135154
builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
@@ -168,14 +187,6 @@ def translate_safe_generator_call(
168187
+ [builder.accept(arg) for arg in expr.args[1:]]),
169188
builder.node_type(expr), expr.line, expr.arg_kinds, expr.arg_names)
170189
else:
171-
if len(expr.args) == 1 and callee.fullname == "builtins.tuple":
172-
val = sequence_from_generator_preallocate_helper(
173-
builder, expr.args[0],
174-
empty_op_llbuilder=builder.builder.new_tuple_with_length,
175-
set_item_op=new_tuple_set_item_op)
176-
if val is not None:
177-
return val
178-
179190
return builder.call_refexpr_with_args(
180191
expr, callee,
181192
([translate_list_comprehension(builder, expr.args[0])]

mypyc/test-data/fixtures/ir.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ def __init__(self, x: object, base: int = 10) -> None: pass
3434
def __add__(self, n: int) -> int: pass
3535
def __sub__(self, n: int) -> int: pass
3636
def __mul__(self, n: int) -> int: pass
37+
def __pow__(self, n: int, modulo: Optional[int] = None) -> int: pass
3738
def __floordiv__(self, x: int) -> int: pass
3839
def __mod__(self, x: int) -> int: pass
3940
def __neg__(self) -> int: pass

mypyc/test-data/run-lists.test

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,3 +354,6 @@ def test() -> None:
354354
source_d = [True, False]
355355
d = [not x for x in source_d]
356356
assert d == [False, True]
357+
source_e = [0, 1, 2]
358+
e = list((x ** 2) for x in (y + 2 for y in source_e))
359+
assert e == [4, 9, 16]

0 commit comments

Comments
 (0)