Optimise iteration over containers declared with a "Dict[...]" type annotation.

See #1979.

CHANGES.rst

@@ -22,6 +22,9 @@ Features added
 
 * Safe integer loops (< range(2^30)) are optimised into C loops.
 
+* Some PEP-484/526 container type declarations are now considered for
+  loop optimisations.
+
 * Python compatible ``cython.*`` types can now be mixed with type declarations
   in Cython syntax.
 

Cython/Compiler/ExprNodes.py

@@ -1863,6 +1863,7 @@ class NameNode(AtomicExprNode):
         if atype is None:
             atype = unspecified_type if as_target and env.directives['infer_types'] != False else py_object_type
         self.entry = env.declare_var(name, atype, self.pos, is_cdef=not as_target)
+        self.entry.annotation = annotation
 
     def analyse_as_module(self, env):
         # Try to interpret this as a reference to a cimported module.

Cython/Compiler/Nodes.py

@@ -2159,7 +2159,10 @@ class FuncDefNode(StatNode, BlockNode):
             error(arg.pos, "Invalid use of 'void'")
         elif not arg.type.is_complete() and not (arg.type.is_array or arg.type.is_memoryviewslice):
             error(arg.pos, "Argument type '%s' is incomplete" % arg.type)
-        return env.declare_arg(arg.name, arg.type, arg.pos)
+        entry = env.declare_arg(arg.name, arg.type, arg.pos)
+        if arg.annotation:
+            entry.annotation = arg.annotation
+        return entry
 
     def generate_arg_type_test(self, arg, code):
         # Generate type test for one argument.
@@ -2775,6 +2778,7 @@ class DefNode(FuncDefNode):
                 name_declarator, type = formal_arg.analyse(scope, nonempty=1)
                 cfunc_args.append(PyrexTypes.CFuncTypeArg(name=name_declarator.name,
                                                           cname=None,
+                                                          annotation=formal_arg.annotation,
                                                           type=py_object_type,
                                                           pos=formal_arg.pos))
             cfunc_type = PyrexTypes.CFuncType(return_type=py_object_type,

Cython/Compiler/Optimize.py

@@ -187,7 +187,18 @@ class IterationTransform(Visitor.EnvTransform):
         return self._optimise_for_loop(node, node.iterator.sequence)
 
     def _optimise_for_loop(self, node, iterator, reversed=False):
-        if iterator.type is Builtin.dict_type:
+        annotation_type = None
+        if (iterator.is_name or iterator.is_attribute) and iterator.entry and iterator.entry.annotation:
+            annotation = iterator.entry.annotation
+            if annotation.is_subscript:
+                annotation = annotation.base  # container base type
+            # FIXME: generalise annotation evaluation => maybe provide a "qualified name" also for imported names?
+            if annotation.is_name and annotation.entry and annotation.entry.qualified_name == 'typing.Dict':
+                annotation_type = Builtin.dict_type
+            elif annotation.is_name and annotation.name == 'Dict':
+                annotation_type = Builtin.dict_type
+
+        if Builtin.dict_type in (iterator.type, annotation_type):
             # like iterating over dict.keys()
             if reversed:
                 # CPython raises an error here: not a sequence

Cython/Compiler/PyrexTypes.py

@@ -3153,15 +3153,18 @@ class CFuncTypeArg(BaseType):
     or_none = False
     accept_none = True
     accept_builtin_subtypes = False
+    annotation = None
 
     subtypes = ['type']
 
-    def __init__(self, name, type, pos, cname=None):
+    def __init__(self, name, type, pos, cname=None, annotation=None):
         self.name = name
         if cname is not None:
             self.cname = cname
         else:
             self.cname = Naming.var_prefix + name
+        if annotation is not None:
+            self.annotation = annotation
         self.type = type
         self.pos = pos
         self.needs_type_test = False # TODO: should these defaults be set in analyse_types()?
@@ -3175,6 +3178,7 @@ class CFuncTypeArg(BaseType):
     def specialize(self, values):
         return CFuncTypeArg(self.name, self.type.specialize(values), self.pos, self.cname)
 
+
 class ToPyStructUtilityCode(object):
 
     requires = None

Cython/Compiler/Symtab.py

@@ -60,6 +60,7 @@ class Entry(object):
     # cname            string     C name of entity
     # type             PyrexType  Type of entity
     # doc              string     Doc string
+    # annotation       ExprNode   PEP 484/526 annotation
     # init             string     Initial value
     # visibility       'private' or 'public' or 'extern'
     # is_builtin       boolean    Is an entry in the Python builtins dict
@@ -138,6 +139,7 @@ class Entry(object):
     inline_func_in_pxd = False
     borrowed = 0
     init = ""
+    annotation = None
     visibility = 'private'
     is_builtin = 0
     is_cglobal = 0

tests/run/pep526_variable_annotations.py

@@ -107,10 +107,57 @@ d['b']: int      # Annotates d['b'] with int.
 (y): int = 0  # Same situation here.
 
 
+@cython.test_assert_path_exists(
+    "//WhileStatNode",
+    "//WhileStatNode//DictIterationNextNode",
+)
+def iter_declared_dict(d):
+    """
+    >>> d = {1.1: 2.5, 3.3: 4.5}
+    >>> iter_declared_dict(d)
+    7.0
+
+    >>> class D(object):
+    ...     def __getitem__(self, x): return 2
+    ...     def __iter__(self): return iter([1, 2, 3])
+    >>> iter_declared_dict(D())
+    6.0
+    """
+    typed_dict : Dict[float, float] = d
+    s = 0.0
+    for key in typed_dict:
+        s += d[key]
+    return s
+
+
+@cython.test_assert_path_exists(
+    "//WhileStatNode",
+    "//WhileStatNode//DictIterationNextNode",
+)
+def iter_declared_dict_arg(d : Dict[float, float]):
+    """
+    >>> d = {1.1: 2.5, 3.3: 4.5}
+    >>> iter_declared_dict_arg(d)
+    7.0
+
+    >>> class D(object):
+    ...     def __getitem__(self, x): return 2
+    ...     def __iter__(self): return iter([1, 2, 3])
+    >>> iter_declared_dict_arg(D())
+    6.0
+    """
+    s = 0.0
+    for key in d:
+        s += d[key]
+    return s
+
 
 _WARNINGS = """
 37:19: Unknown type declaration in annotation, ignoring
 38:12: Unknown type declaration in annotation, ignoring
 39:18: Unknown type declaration in annotation, ignoring
 73:19: Unknown type declaration in annotation, ignoring
+# FIXME: these are sort-of evaluated now, so the warning is misleading
+126:21: Unknown type declaration in annotation, ignoring
+137:35: Unknown type declaration in annotation, ignoring
 """