Commit 415c0496 authored by Max Bachmann's avatar Max Bachmann Committed by GitHub

Optimize enumerate() with start argument and (GH-4077)

Also fixes the control flow analysis for loop variable assignments in this case.
Closes https://github.com/cython/cython/issues/3787
parent cf863062
......@@ -980,7 +980,7 @@ class ControlFlowAnalysis(CythonTransform):
if not entry or entry.is_builtin:
if function.name == 'reversed' and len(sequence.args) == 1:
sequence = sequence.args[0]
elif function.name == 'enumerate' and len(sequence.args) == 1:
elif function.name == 'enumerate' and len(sequence.args) in (1, 2):
if target.is_sequence_constructor and len(target.args) == 2:
iterator = sequence.args[0]
if iterator.is_name:
......@@ -1016,7 +1016,11 @@ class ControlFlowAnalysis(CythonTransform):
# Python strings, etc., while correctly falling back to an
# object type when the base type cannot be handled.
self.mark_assignment(target, node.item)
self.mark_assignment(target, ExprNodes.IndexNode(
node.pos,
base=sequence,
index=ExprNodes.IntNode(target.pos, value='PY_SSIZE_T_MAX',
type=PyrexTypes.c_py_ssize_t_type)))
def visit_AsyncForStatNode(self, node):
return self.visit_ForInStatNode(node)
......
......@@ -111,7 +111,7 @@ class MarkParallelAssignments(EnvTransform):
if not entry or entry.is_builtin:
if function.name == 'reversed' and len(sequence.args) == 1:
sequence = sequence.args[0]
elif function.name == 'enumerate' and len(sequence.args) == 1:
elif function.name == 'enumerate' and len(sequence.args) in (1, 2):
if target.is_sequence_constructor and len(target.args) == 2:
iterator = sequence.args[0]
if iterator.is_name:
......
from __future__ import print_function
cimport cython
from cython cimport typeof
bytes_abc = b'abc'
bytes_ABC = b'ABC'
......@@ -224,3 +226,27 @@ def for_pyucs4_in_enumerate_unicode(unicode s):
return i
else:
return 'X'
@cython.test_assert_path_exists("//ForFromStatNode")
@cython.test_fail_if_path_exists("//ForInStatNode")
def for_infered_type_in_enumerate_unicode(unicode s):
"""
>>> for_infered_type_in_enumerate_unicode(unicode_abc)
Py_ssize_t Py_UCS4
'X'
>>> for_infered_type_in_enumerate_unicode(unicode_ABC)
Py_ssize_t Py_UCS4
2
>>> for_infered_type_in_enumerate_unicode(unicode_abc_null)
Py_ssize_t Py_UCS4
'X'
>>> for_infered_type_in_enumerate_unicode(unicode_ABC_null)
Py_ssize_t Py_UCS4
4
"""
print(typeof(i), typeof(c))
for i, c in enumerate(s):
if c == u'C':
return i
else:
return 'X'
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment