Improved REMOVE object as a "falsey" + Parallel visitor supports editing

Related GraphQL-js commit: graphql/graphql-js@bd87fd3

Author: syrusakbary

graphql/core/language/visitor.py

@@ -5,8 +5,17 @@
 from . import ast
 from .visitor_meta import QUERY_DOCUMENT_KEYS, VisitorMeta
 
+
+class Falsey(object):
+    def __nonzero__(self):
+        return False
+
+    def __bool__(self):
+        return False
+
+
 BREAK = object()
-REMOVE = object()
+REMOVE = Falsey()
 
 
 class Stack(object):
@@ -90,7 +99,7 @@ def visit(root, visitor, key_map=None):
                 key = None
                 node = new_root
 
-            if node is None:
+            if node is REMOVE or node is None:
                 continue
 
             if parent:
@@ -179,15 +188,19 @@ def enter(self, node, key, parent, path, ancestors):
                     self.skipping[i] = node
                 elif result is BREAK:
                     self.skipping[i] = BREAK
+                elif result is not None:
+                    return result
 
     def leave(self, node, key, parent, path, ancestors):
         for i, visitor in enumerate(self.visitors):
             if not self.skipping[i]:
                 result = visitor.leave(node, key, parent, path, ancestors)
                 if result is BREAK:
                     self.skipping[i] = BREAK
+                elif result is not None and result is not False:
+                    return result
             elif self.skipping[i] == node:
-                self.skipping[i] = None
+                self.skipping[i] = REMOVE
 
 
 class TypeInfoVisitor(Visitor):

tests/core_language/test_visitor.py

@@ -674,3 +674,115 @@ def leave(self, node, key, parent, *args):
         ['break-b', 'leave', 'SelectionSet', None],
         ['break-b', 'leave', 'Field', None]
     ]
+
+
+def test_visits_in_pararell_allows_for_editing_on_enter():
+    visited = []
+    ast = parse('{ a, b, c { a, b, c } }', no_location=True)
+
+    class TestVisitor1(Visitor):
+
+        def enter(self, node, key, parent, *args):
+            if type(node).__name__ == 'Field' and node.name.value == 'b':
+                return REMOVE
+
+    class TestVisitor2(Visitor):
+
+        def enter(self, node, key, parent, *args):
+            visited.append(
+                ['enter', type(node).__name__, getattr(node, 'value', None)])
+
+        def leave(self, node, key, parent, *args):
+            visited.append(
+                ['leave', type(node).__name__, getattr(node, 'value', None)])
+
+    edited_ast = visit(ast, ParallelVisitor([TestVisitor1(), TestVisitor2()]))
+
+    assert ast == parse('{ a, b, c { a, b, c } }', no_location=True)
+    assert edited_ast == parse('{ a,    c { a,    c } }', no_location=True)
+
+    assert visited == [
+        ['enter', 'Document', None],
+        ['enter', 'OperationDefinition', None],
+        ['enter', 'SelectionSet', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'a'],
+        ['leave', 'Name', 'a'],
+        ['leave', 'Field', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'c'],
+        ['leave', 'Name', 'c'],
+        ['enter', 'SelectionSet', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'a'],
+        ['leave', 'Name', 'a'],
+        ['leave', 'Field', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'c'],
+        ['leave', 'Name', 'c'],
+        ['leave', 'Field', None],
+        ['leave', 'SelectionSet', None],
+        ['leave', 'Field', None],
+        ['leave', 'SelectionSet', None],
+        ['leave', 'OperationDefinition', None],
+        ['leave', 'Document', None]
+    ]
+
+
+def test_visits_in_pararell_allows_for_editing_on_leave():
+    visited = []
+    ast = parse('{ a, b, c { a, b, c } }', no_location=True)
+
+    class TestVisitor1(Visitor):
+
+        def leave(self, node, key, parent, *args):
+            if type(node).__name__ == 'Field' and node.name.value == 'b':
+                return REMOVE
+
+    class TestVisitor2(Visitor):
+
+        def enter(self, node, key, parent, *args):
+            visited.append(
+                ['enter', type(node).__name__, getattr(node, 'value', None)])
+
+        def leave(self, node, key, parent, *args):
+            visited.append(
+                ['leave', type(node).__name__, getattr(node, 'value', None)])
+
+    edited_ast = visit(ast, ParallelVisitor([TestVisitor1(), TestVisitor2()]))
+
+    assert ast == parse('{ a, b, c { a, b, c } }', no_location=True)
+    assert edited_ast == parse('{ a,    c { a,    c } }', no_location=True)
+
+    assert visited == [
+        ['enter', 'Document', None],
+        ['enter', 'OperationDefinition', None],
+        ['enter', 'SelectionSet', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'a'],
+        ['leave', 'Name', 'a'],
+        ['leave', 'Field', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'b'],
+        ['leave', 'Name', 'b'],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'c'],
+        ['leave', 'Name', 'c'],
+        ['enter', 'SelectionSet', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'a'],
+        ['leave', 'Name', 'a'],
+        ['leave', 'Field', None],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'b'],
+        ['leave', 'Name', 'b'],
+        ['enter', 'Field', None],
+        ['enter', 'Name', 'c'],
+        ['leave', 'Name', 'c'],
+        ['leave', 'Field', None],
+        ['leave', 'SelectionSet', None],
+        ['leave', 'Field', None],
+        ['leave', 'SelectionSet', None],
+        ['leave', 'OperationDefinition', None],
+        ['leave', 'Document', None]
+    ]