Comprehensions

pyproject.toml

@@ -1,4 +1,5 @@
 [tool.black]
 
 target-version = ['py27', 'py33', 'py34', 'py35', 'py36', 'py37', 'py38', 'py39']
-exclude = 'ropetest'
+include = 'rope/.*\.pyi?$'
+force-exclude = 'ropetest'

rope/refactor/extract.py

@@ -843,6 +843,31 @@ def _AugAssign(self, node):
     def _ClassDef(self, node):
         self._written_variable(node.name, node.lineno)
 
+    def _ListComp(self, node):
+        self._comp_exp(node)
+
+    def _GeneratorExp(self, node):
+        self._comp_exp(node)
+
+    def _SetComp(self, node):
+        self._comp_exp(node)
+
+    def _DictComp(self, node):
+        self._comp_exp(node)
+
+    def _comp_exp(self, node):
+        read = OrderedSet(self.read)
+        written = OrderedSet(self.written)
+        maybe_written = OrderedSet(self.maybe_written)
+
+        for child in ast.get_child_nodes(node):
+            ast.walk(child, self)
+
+        comp_names = set([name.target.id for name in node.generators])
+        self.read = self.read - comp_names | read
+        self.written = self.written - comp_names | written
+        self.maybe_written = self.maybe_written - comp_names | maybe_written
+
     def _If(self, node):
         self._handle_conditional_node(node)
 

ropetest/refactor/extracttest.py

@@ -1395,7 +1395,7 @@ def test_extract_function_with_for_else_statemant(self):
         self.assertEqual(expected, refactored)
 
     def test_extract_function_with_for_else_statemant_more(self):
-        """TODO: fixed code to test passed """
+        """TODO: fixed code to test passed"""
         code = (
             "def a_func():\n"
             "    for i in range(10):\n"
@@ -1979,6 +1979,134 @@ def second_method(someargs):
 
         self.assertEqual(expected, refactored)
 
+    def test_extract_with_list_comprehension(self):
+        code = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = sum([x for x in y])
+                b = sum([x for x in y])
+
+                print(a, b)
+
+            f()
+        """)
+        extract_target = "    a = sum([x for x in y])\n"
+        start, end = code.index(extract_target), code.index(extract_target) + len(
+            extract_target
+        )
+        refactored = self.do_extract_method(code, start, end, "_a")
+        expected = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = _a(y)
+                b = sum([x for x in y])
+
+                print(a, b)
+
+            def _a(y):
+                a = sum([x for x in y])
+                return a
+
+            f()
+        """)
+        self.assertEqual(expected, refactored)
+
+    def test_extract_with_generator(self):
+        code = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = sum(x for x in y)
+                b = sum(x for x in y)
+
+                print(a, b)
+
+            f()
+        """)
+        extract_target = "    a = sum(x for x in y)\n"
+        start, end = code.index(extract_target), code.index(extract_target) + len(
+            extract_target
+        )
+        refactored = self.do_extract_method(code, start, end, "_a")
+        expected = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = _a(y)
+                b = sum(x for x in y)
+
+                print(a, b)
+
+            def _a(y):
+                a = sum(x for x in y)
+                return a
+
+            f()
+        """)
+        self.assertEqual(expected, refactored)
+
+    def test_extract_with_set_comprehension(self):
+        code = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = sum({x for x in y})
+                b = sum({x for x in y})
+
+                print(a, b)
+
+            f()
+        """)
+        extract_target = "    a = sum({x for x in y})\n"
+        start, end = code.index(extract_target), code.index(extract_target) + len(
+            extract_target
+        )
+        refactored = self.do_extract_method(code, start, end, "_a")
+        expected = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = _a(y)
+                b = sum({x for x in y})
+
+                print(a, b)
+
+            def _a(y):
+                a = sum({x for x in y})
+                return a
+
+            f()
+        """)
+        self.assertEqual(expected, refactored)
+
+    def test_extract_with_dict_comprehension(self):
+        code = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = sum({x: x for x in y})
+                b = sum({x: x for x in y})
+
+                print(a, b)
+
+            f()
+        """)
+        extract_target = "    a = sum({x: x for x in y})\n"
+        start, end = code.index(extract_target), code.index(extract_target) + len(
+            extract_target
+        )
+        refactored = self.do_extract_method(code, start, end, "_a")
+        expected = dedent("""\
+            def f():
+                y = [1,2,3,4]
+                a = _a(y)
+                b = sum({x: x for x in y})
+
+                print(a, b)
+
+            def _a(y):
+                a = sum({x: x for x in y})
+                return a
+
+            f()
+        """)
+        self.assertEqual(expected, refactored)
+
     def test_extract_function_expression_with_assignment_to_attribute(self):
         code = dedent("""\
             class A(object):