From 39ca71e4926b71e0fd8d32127bce81ecb961f0a2 Mon Sep 17 00:00:00 2001
From: Stefan Behnel <scoder@users.berlios.de>
Date: Sun, 7 Nov 2010 18:29:24 +0100
Subject: [PATCH] optimise iteration over dict.keys/values/items() in -3 mode

---
 Cython/Compiler/Optimize.py |  8 +++++---
 tests/run/cython3.pyx       | 25 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/Cython/Compiler/Optimize.py b/Cython/Compiler/Optimize.py
index ba96caf3..c1ca6173 100644
--- a/Cython/Compiler/Optimize.py
+++ b/Cython/Compiler/Optimize.py
@@ -73,6 +73,7 @@ class IterationTransform(Visitor.VisitorTransform):
 
     def visit_ModuleNode(self, node):
         self.current_scope = node.scope
+        self.module_scope = node.scope
         self.visitchildren(node)
         return node
 
@@ -168,12 +169,13 @@ class IterationTransform(Visitor.VisitorTransform):
             dict_obj = function.obj
             method = function.attribute
 
+            is_py3 = self.module_scope.context.language_level >= 3
             keys = values = False
-            if method == 'iterkeys':
+            if method == 'iterkeys' or (is_py3 and method == 'keys'):
                 keys = True
-            elif method == 'itervalues':
+            elif method == 'itervalues' or (is_py3 and method == 'values'):
                 values = True
-            elif method == 'iteritems':
+            elif method == 'iteritems' or (is_py3 and method == 'items'):
                 keys = values = True
             else:
                 return node
diff --git a/tests/run/cython3.pyx b/tests/run/cython3.pyx
index b197eb35..0c59349c 100644
--- a/tests/run/cython3.pyx
+++ b/tests/run/cython3.pyx
@@ -1,5 +1,7 @@
 # cython: language_level=3
 
+cimport cython
+
 try:
     sorted
 except NameError:
@@ -74,3 +76,26 @@ def dict_comp():
     result = {x:x*2 for x in range(5) if x % 2 == 0}
     assert x == 'abc' # don't leak
     return result
+
+# in Python 3, d.keys/values/items() are the iteration methods
+@cython.test_assert_path_exists(
+    "//WhileStatNode",
+    "//WhileStatNode/SimpleCallNode",
+    "//WhileStatNode/SimpleCallNode/NameNode")
+@cython.test_fail_if_path_exists(
+    "//ForInStatNode")
+def dict_iter(dict d):
+    """
+    >>> d = {'a' : 1, 'b' : 2, 'c' : 3}
+    >>> keys, values, items = dict_iter(d)
+    >>> sorted(keys)
+    ['a', 'b', 'c']
+    >>> sorted(values)
+    [1, 2, 3]
+    >>> sorted(items)
+    [('a', 1), ('b', 2), ('c', 3)]
+    """
+    keys = [ key for key in d.keys() ]
+    values = [ value for value in d.values() ]
+    items = [ item for item in d.items() ]
+    return keys, values, items
-- 
2.26.2