Fix locals() at module and class scope, see #731

Cython/Compiler/ExprNodes.py

@@ -5455,7 +5455,7 @@ class GlobalsExprNode(AtomicExprNode):
         code.put_gotref(self.result())
 
 
-class LocalsExprNode(DictNode):
+class FuncLocalsExprNode(DictNode):
     def __init__(self, pos, env):
         local_vars = [var.name for var in env.entries.values() if var.name]
         items = [DictItemNode(pos, key=IdentifierStringNode(pos, value=var),
@@ -5464,6 +5464,31 @@ class LocalsExprNode(DictNode):
         DictNode.__init__(self, pos, key_value_pairs=items,
                           exclude_null_values=True)
 
+
+class PyClassLocalsExprNode(AtomicExprNode):
+    def __init__(self, pos, pyclass_dict):
+        AtomicExprNode.__init__(self, pos)
+        self.pyclass_dict = pyclass_dict
+
+    def analyse_types(self, env):
+        self.type = self.pyclass_dict.type
+        self.is_tmep = 0
+
+    def result(self):
+        return self.pyclass_dict.result()
+
+    def generate_result_code(self, code):
+        pass
+
+
+def LocalsExprNode(pos, scope_node, env):
+    if env.is_module_scope:
+        return GlobalsExprNode(pos)
+    if env.is_py_class_scope:
+        return PyClassLocalsExprNode(pos, scope_node.dict)
+    return FuncLocalsExprNode(pos, env)
+
+
 #-------------------------------------------------------------------
 #
 #  Unary operator nodes

Cython/Compiler/ParseTreeTransforms.py

@@ -2269,7 +2269,7 @@ class TransformBuiltinMethods(EnvTransform):
                           % len(node.args))
                 if len(node.args) > 0:
                     return node # nothing to do
-            return ExprNodes.LocalsExprNode(pos, lenv)
+            return ExprNodes.LocalsExprNode(pos, self.current_scope_node(), lenv)
         else: # dir()
             if len(node.args) > 1:
                 error(self.pos, "Builtin 'dir()' called with wrong number of args, expected 0-1, got %d"
@@ -2277,6 +2277,17 @@ class TransformBuiltinMethods(EnvTransform):
             if len(node.args) > 0:
                 # optimised in Builtin.py
                 return node
+            if lenv.is_py_class_scope or lenv.is_module_scope:
+                if lenv.is_py_class_scope:
+                    pyclass = self.current_scope_node()
+                    locals_dict = ExprNodes.CloneNode(pyclass.dict)
+                else:
+                    locals_dict = ExprNodes.GlobalsExprNode(pos)
+                return ExprNodes.SimpleCallNode(
+                    pos,
+                    function=ExprNodes.AttributeNode(
+                        pos, obj=locals_dict, attribute="keys"),
+                    args=[])
             local_names = [ var.name for var in lenv.entries.values() if var.name ]
             items = [ ExprNodes.IdentifierStringNode(pos, value=var)
                       for var in local_names ]

tests/run/locals_T732.pyx

+# mode: run
+# ticket: 731
+# tags: locals, vars, dir
+
+LOCALS = locals()
+GLOBALS = globals()
+DIR_SAME = sorted(dir()) == sorted(globals().keys())
+
+def test_module_locals_and_dir():
+    """
+    >>> LOCALS is GLOBALS
+    True
+    >>> DIR_SAME
+    True
+    """
+
+def test_class_locals_and_dir():
+    """
+    >>> klass = test_class_locals_and_dir()
+    >>> 'visible' in klass.locs and 'not_visible' not in klass.locs
+    True
+    >>> klass.names
+    ['visible']
+    """
+    not_visible = 1234
+    class Foo:
+        visible = 4321
+        names = dir()
+        locs = locals()
+    return Foo