add support for DictComp and SetComp to patched_ast

python-rope · Jul 17, 2015 · dc1cb01 · dc1cb01
1 parent d8b4f98
commit dc1cb01
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 0 deletions.
diff --git a/rope/refactor/patchedast.py b/rope/refactor/patchedast.py
@@ -493,6 +493,19 @@ def _ListComp(self, node):
         children.append(']')
         self._handle(node, children)
 
+    def _SetComp(self, node):
+        children = ['{', node.elt]
+        children.extend(node.generators)
+        children.append('}')
+        self._handle(node, children)
+
+    def _DictComp(self, node):
+        children = ['{']
+        children.extend([node.key, ':', node.value])
+        children.extend(node.generators)
+        children.append('}')
+        self._handle(node, children)
+
     def _Module(self, node):
         self._handle(node, list(node.body), eat_spaces=True)
 

diff --git a/ropetest/refactor/patchedasttest.py b/ropetest/refactor/patchedasttest.py
@@ -578,6 +578,43 @@ def test_list_comp_node_with_multiple_comprehensions(self):
             'comprehension', ['for', ' ', 'Name', ' ', 'in', ' ',
                               'Call', ' ', 'if', ' ', 'Name'])
 
+    def test_set_comp_node(self):
+        # make sure we are in a python version with set comprehensions
+        source = '{i for i in range(1) if True}\n'
+
+        try:
+            eval(source)
+        except SyntaxError:
+            return
+
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+        checker.check_region('SetComp', 0, len(source) - 1)
+        checker.check_children(
+            'SetComp', ['{', '', 'Name', ' ', 'comprehension', '', '}'])
+        checker.check_children(
+            'comprehension', ['for', ' ', 'Name', ' ', 'in', ' ',
+                              'Call', ' ', 'if', ' ', 'Name'])
+
+    def test_dict_comp_node(self):
+        # make sure we are in a python version with dict comprehensions
+        source = '{i:i for i in range(1) if True}\n'
+
+        try:
+            eval(source)
+        except SyntaxError:
+            return
+
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+        checker.check_region('DictComp', 0, len(source) - 1)
+        checker.check_children(
+            'DictComp', ['{', '', 'Name', '', ':', '', 'Name',
+                         ' ', 'comprehension', '', '}'])
+        checker.check_children(
+            'comprehension', ['for', ' ', 'Name', ' ', 'in', ' ',
+                              'Call', ' ', 'if', ' ', 'Name'])
+
     def test_ext_slice_node(self):
         source = 'x = xs[0,:]\n'
         ast_frag = patchedast.get_patched_ast(source, True)