Comprhensions with scope

CHANGELOG.md

@@ -7,7 +7,7 @@
 ## Bug fixes
 - #391, #396 Extract method similar no longer replace the left-hand side of assignment
 - #303 Fix inlining into f-string containing quote characters
-
+- Added scopes for comprehension expressions as part of #293
 
 
 

rope/base/pyobjects.py

@@ -247,6 +247,10 @@ class PyFunction(PyDefinedObject, AbstractFunction):
     """Only a placeholder"""
 
 
+class PyComprehension(PyDefinedObject, PyObject):
+    """Only a placeholder"""
+
+
 class PyClass(PyDefinedObject, AbstractClass):
     """Only a placeholder"""
 

rope/base/pyobjectsdef.py

@@ -1,3 +1,4 @@
+from rope.base.pynames import DefinedName
 import rope.base.builtins
 import rope.base.codeanalyze
 import rope.base.evaluate
@@ -121,6 +122,18 @@ def decorators(self):
             return getattr(self.ast_node, "decorators", None)
 
 
+class PyComprehension(pyobjects.PyComprehension):
+    def __init__(self, pycore, ast_node, parent):
+        self.visitor_class = _ComprehensionVisitor
+        rope.base.pyobjects.PyObject.__init__(self, type_="Comp")
+        rope.base.pyobjects.PyDefinedObject.__init__(self, pycore, ast_node, parent)
+
+    def _create_scope(self):
+        return rope.base.pyscopes.ComprehensionScope(
+            self.pycore, self, _ComprehensionVisitor
+        )
+
+
 class PyClass(pyobjects.PyClass):
     def __init__(self, pycore, ast_node, parent):
         self.visitor_class = _ClassVisitor
@@ -329,23 +342,16 @@ class _ExpressionVisitor(object):
     def __init__(self, scope_visitor):
         self.scope_visitor = scope_visitor
 
-    def _assigned(self, name, assignment=None):
-        self.scope_visitor._assigned(name, assignment)
-
     def _GeneratorExp(self, node):
-        for child in ["elt", "key", "value"]:
-            if hasattr(node, child):
-                ast.walk(getattr(node, child), self)
-        for comp in node.generators:
-            ast.walk(comp.target, _AssignVisitor(self))
-            ast.walk(comp, self)
-            for if_ in comp.ifs:
-                ast.walk(if_, self)
+        list_comp = PyComprehension(
+            self.scope_visitor.pycore, node, self.scope_visitor.owner_object
+        )
+        self.scope_visitor.defineds.append(list_comp)
 
-    def _ListComp(self, node):
+    def _SetComp(self, node):
         self._GeneratorExp(node)
 
-    def _SetComp(self, node):
+    def _ListComp(self, node):
         self._GeneratorExp(node)
 
     def _DictComp(self, node):
@@ -396,6 +402,7 @@ def _Slice(self, node):
 
 class _ScopeVisitor(_ExpressionVisitor):
     def __init__(self, pycore, owner_object):
+        _ExpressionVisitor.__init__(self, scope_visitor=self)
         self.pycore = pycore
         self.owner_object = owner_object
         self.names = {}
@@ -562,6 +569,19 @@ def _Global(self, node):
             self.names[name] = pyname
 
 
+class _ComprehensionVisitor(_ScopeVisitor):
+    def _comprehension(self, node):
+        ast.walk(node.target, self)
+        ast.walk(node.iter, self)
+
+    def _Name(self, node):
+        if isinstance(node.ctx, ast.Store):
+            self.names[node.id] = DefinedName(self._get_pyobject(node))
+
+    def _get_pyobject(self, node):
+        return pyobjects.PyDefinedObject(None, node, self.owner_object)
+
+
 class _GlobalVisitor(_ScopeVisitor):
     def __init__(self, pycore, owner_object):
         super(_GlobalVisitor, self).__init__(pycore, owner_object)

rope/base/pyscopes.py

@@ -143,6 +143,42 @@ def builtin_names(self):
         return rope.base.builtins.builtins.get_attributes()
 
 
+class ComprehensionScope(Scope):
+    def __init__(self, pycore, pyobject, visitor):
+        super(ComprehensionScope, self).__init__(
+            pycore, pyobject, pyobject.parent.get_scope()
+        )
+        self.names = None
+        self.returned_asts = None
+        self.defineds = None
+        self.visitor = visitor
+
+    def _get_names(self):
+        if self.names is None:
+            self._visit_comprehension()
+        return self.names
+
+    def get_names(self):
+        return self._get_names()
+
+    def _visit_comprehension(self):
+        if self.names is None:
+            new_visitor = self.visitor(self.pycore, self.pyobject)
+            for node in ast.get_child_nodes(self.pyobject.get_ast()):
+                ast.walk(node, new_visitor)
+            self.names = new_visitor.names
+            self.names.update(self.parent.get_names())
+            self.defineds = new_visitor.defineds
+
+    def get_logical_end(self):
+        return self.get_start()
+
+    logical_end = property(get_logical_end)
+
+    def get_body_start(self):
+        return self.get_start()
+
+
 class FunctionScope(Scope):
     def __init__(self, pycore, pyobject, visitor):
         super(FunctionScope, self).__init__(

rope/base/utils/__init__.py

@@ -136,7 +136,7 @@ def is_inline_body():
         body_col_offset = node.body[0].col_offset
         return indent_col_offset < body_col_offset
 
-    if sys.version_info >= (3, 8):
+    if sys.version_info >= (3, 8) or not hasattr(node, "body"):
         return node.lineno
 
     possible_def_line = (

ropetest/pyscopestest.py

@@ -52,15 +52,19 @@ def test_list_comprehension_scope_inside_assignment(self):
         )
         self.assertEqual(
             list(sorted(scope.get_defined_names())),
-            ["a_var", "b_var", "c_var"],
+            ["a_var"],
+        )
+        self.assertEqual(
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
+            ["b_var", "c_var"],
         )
 
     def test_list_comprehension_scope(self):
         scope = libutils.get_string_scope(
             self.project, "[b_var + d_var for b_var, c_var in e_var]\n"
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
             ["b_var", "c_var"],
         )
 
@@ -69,7 +73,7 @@ def test_set_comprehension_scope(self):
             self.project, "{b_var + d_var for b_var, c_var in e_var}\n"
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
             ["b_var", "c_var"],
         )
 
@@ -78,7 +82,7 @@ def test_generator_comprehension_scope(self):
             self.project, "(b_var + d_var for b_var, c_var in e_var)\n"
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
             ["b_var", "c_var"],
         )
 
@@ -87,7 +91,7 @@ def test_dict_comprehension_scope(self):
             self.project, "{b_var: d_var for b_var, c_var in e_var}\n"
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
             ["b_var", "c_var"],
         )
 
@@ -102,8 +106,12 @@ def test_inline_assignment_in_comprehensions(self):
             ]""",
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
-            ["a_var", "b_var", "f_var", "h_var", "i_var", "j_var"],
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
+            ["a_var", "b_var", "f_var"],
+        )
+        self.assertEqual(
+            list(sorted(scope.get_scopes()[0].get_scopes()[0].get_defined_names())),
+            ["i_var", "j_var"],
         )
 
     def test_nested_comprehension(self):
@@ -116,8 +124,12 @@ def test_nested_comprehension(self):
             ]\n""",
         )
         self.assertEqual(
-            list(sorted(scope.get_defined_names())),
-            ["b_var", "c_var", "e_var"],
+            list(sorted(scope.get_scopes()[0].get_defined_names())),
+            ["b_var", "c_var"],
+        )
+        self.assertEqual(
+            list(sorted(scope.get_scopes()[0].get_scopes()[0].get_defined_names())),
+            ["e_var"],
         )
 
     def test_simple_class_scope(self):
@@ -330,3 +342,56 @@ def test_getting_defined_names_for_modules(self):
         self.assertTrue("A" in scope.get_names())
         self.assertTrue("open" not in scope.get_defined_names())
         self.assertTrue("A" in scope.get_defined_names())
+
+    def test_get_inner_scope_for_list_comprhension_with_many_targets(self):
+        scope = libutils.get_string_scope(
+            self.project, "a = [(i, j) for i,j in enumerate(range(10))]\n"
+        )
+        self.assertEqual(len(scope.get_scopes()), 1)
+        self.assertNotIn("i", scope)
+        self.assertNotIn("j", scope)
+        self.assertIn("i", scope.get_scopes()[0])
+        self.assertIn("j", scope.get_scopes()[0])
+
+    def test_get_inner_scope_for_generator(self):
+        scope = libutils.get_string_scope(self.project, "a = (i for i in range(10))\n")
+        self.assertEqual(len(scope.get_scopes()), 1)
+        self.assertNotIn("i", scope)
+        self.assertIn("i", scope.get_scopes()[0])
+
+    def test_get_inner_scope_for_set_comprehension(self):
+        scope = libutils.get_string_scope(self.project, "a = {i for i in range(10)}\n")
+        self.assertEqual(len(scope.get_scopes()), 1)
+        self.assertNotIn("i", scope)
+        self.assertIn("i", scope.get_scopes()[0])
+
+    def test_get_inner_scope_for_dict_comprehension(self):
+        scope = libutils.get_string_scope(
+            self.project, "a = {i:i for i in range(10)}\n"
+        )
+        self.assertEqual(len(scope.get_scopes()), 1)
+        self.assertNotIn("i", scope)
+        self.assertIn("i", scope.get_scopes()[0])
+
+    def test_get_inner_scope_for_nested_list_comprhension(self):
+        scope = libutils.get_string_scope(
+            self.project, "a = [[i + j for j in range(10)] for i in range(10)]\n"
+        )
+
+        self.assertEqual(len(scope.get_scopes()), 1)
+        self.assertNotIn("i", scope)
+        self.assertNotIn("j", scope)
+        self.assertIn("i", scope.get_scopes()[0])
+        self.assertEqual(len(scope.get_scopes()[0].get_scopes()), 1)
+        self.assertIn("j", scope.get_scopes()[0].get_scopes()[0])
+        self.assertIn("i", scope.get_scopes()[0].get_scopes()[0])
+
+
+def suite():
+    result = unittest.TestSuite()
+    result.addTests(unittest.makeSuite(PyCoreScopesTest))
+    return result
+
+
+if __name__ == "__main__":
+    unittest.main()