Add enum type for visitor return values (#96)

Author: Cito

docs/modules/language.rst

@@ -103,17 +103,17 @@ The module also exports the following special symbols which can be used as
 return values in the :class:`Visitor` methods to signal particular actions:
 
 .. data:: BREAK
-   :annotation: = True
+   :annotation: (same as ``True``)
 
    This return value signals that no further nodes shall be visited.
 
 .. data:: SKIP
-   :annotation: = False
+   :annotation: (same as ``False``)
 
    This return value signals that the current node shall be skipped.
 
 .. data:: REMOVE
-   :annotation: = Ellipsis
+   :annotation: (same as``Ellipsis``)
 
    This return value signals that the current node shall be deleted.
 

src/graphql/__init__.py

@@ -181,6 +181,7 @@
     visit,
     ParallelVisitor,
     Visitor,
+    VisitorAction,
     BREAK,
     SKIP,
     REMOVE,
@@ -532,6 +533,7 @@
     "ParallelVisitor",
     "TypeInfoVisitor",
     "Visitor",
+    "VisitorAction",
     "BREAK",
     "SKIP",
     "REMOVE",

src/graphql/language/__init__.py

@@ -22,6 +22,7 @@
     visit,
     Visitor,
     ParallelVisitor,
+    VisitorAction,
     BREAK,
     SKIP,
     REMOVE,
@@ -115,6 +116,7 @@
     "visit",
     "Visitor",
     "ParallelVisitor",
+    "VisitorAction",
     "BREAK",
     "SKIP",
     "REMOVE",

src/graphql/language/visitor.py

@@ -1,4 +1,5 @@
 from copy import copy
+from enum import Enum
 from typing import (
     Any,
     Callable,
@@ -19,6 +20,7 @@
 __all__ = [
     "Visitor",
     "ParallelVisitor",
+    "VisitorAction",
     "visit",
     "BREAK",
     "SKIP",
@@ -28,10 +30,26 @@
 ]
 
 
-# Special return values for the visitor methods:
+class VisitorActionEnum(Enum):
+    """Special return values for the visitor methods.
+
+    You can also use the values of this enum directly.
+    """
+
+    BREAK = True
+    SKIP = False
+    REMOVE = Ellipsis
+
+
+VisitorAction = Optional[VisitorActionEnum]
+
 # Note that in GraphQL.js these are defined differently:
 # BREAK = {}, SKIP = false, REMOVE = null, IDLE = undefined
-BREAK, SKIP, REMOVE, IDLE = True, False, Ellipsis, None
+
+BREAK = VisitorActionEnum.BREAK
+SKIP = VisitorActionEnum.SKIP
+REMOVE = VisitorActionEnum.REMOVE
+IDLE = None
 
 # Default map from visitor kinds to their traversable node attributes:
 QUERY_DOCUMENT_KEYS: Dict[str, Tuple[str, ...]] = {
@@ -253,7 +271,7 @@ def visit(root: Node, visitor: Visitor, visitor_keys=None) -> Any:
             for edit_key, edit_value in edits:
                 if in_array:
                     edit_key -= edit_offset
-                if in_array and edit_value is REMOVE:
+                if in_array and (edit_value is REMOVE or edit_value is Ellipsis):
                     node.pop(edit_key)
                     edit_offset += 1
                 else:
@@ -292,10 +310,10 @@ def visit(root: Node, visitor: Visitor, visitor_keys=None) -> Any:
             if visit_fn:
                 result = visit_fn(visitor, node, key, parent, path, ancestors)
 
-                if result is BREAK:
+                if result is BREAK or result is True:
                     break
 
-                if result is SKIP:
+                if result is SKIP or result is False:
                     if not is_leaving:
                         path_pop()
                         continue
@@ -356,9 +374,9 @@ def enter(self, node, *args):
                 fn = visitor.get_visit_fn(node.kind)
                 if fn:
                     result = fn(visitor, node, *args)
-                    if result is SKIP:
+                    if result is SKIP or result is False:
                         skipping[i] = node
-                    elif result == BREAK:
+                    elif result is BREAK or result is True:
                         skipping[i] = BREAK
                     elif result is not None:
                         return result
@@ -370,9 +388,13 @@ def leave(self, node, *args):
                 fn = visitor.get_visit_fn(node.kind, is_leaving=True)
                 if fn:
                     result = fn(visitor, node, *args)
-                    if result == BREAK:
+                    if result is BREAK or result is True:
                         skipping[i] = BREAK
-                    elif result is not None and result is not SKIP:
+                    elif (
+                        result is not None
+                        and result is not SKIP
+                        and result is not False
+                    ):
                         return result
             elif skipping[i] is node:
                 skipping[i] = None

src/graphql/validation/rules/executable_definitions.py

@@ -8,6 +8,8 @@
     SchemaDefinitionNode,
     SchemaExtensionNode,
     TypeDefinitionNode,
+    VisitorAction,
+    SKIP,
 )
 from . import ASTValidationRule
 
@@ -21,7 +23,7 @@ class ExecutableDefinitionsRule(ASTValidationRule):
     operation or fragment definitions.
     """
 
-    def enter_document(self, node: DocumentNode, *_args):
+    def enter_document(self, node: DocumentNode, *_args) -> VisitorAction:
         for definition in node.definitions:
             if not isinstance(definition, ExecutableDefinitionNode):
                 def_name = (
@@ -41,4 +43,4 @@ def enter_document(self, node: DocumentNode, *_args):
                         f"The {def_name} definition is not executable.", definition,
                     )
                 )
-        return self.SKIP
+        return SKIP

src/graphql/validation/rules/fields_on_correct_type.py

@@ -27,7 +27,7 @@ class FieldsOnCorrectTypeRule(ValidationRule):
     type, or are an allowed meta field such as ``__typename``.
     """
 
-    def enter_field(self, node: FieldNode, *_args):
+    def enter_field(self, node: FieldNode, *_args) -> None:
         type_ = self.context.get_parent_type()
         if not type_:
             return

src/graphql/validation/rules/fragments_on_composite_types.py

@@ -1,5 +1,9 @@
 from ...error import GraphQLError
-from ...language import FragmentDefinitionNode, InlineFragmentNode, print_ast
+from ...language import (
+    FragmentDefinitionNode,
+    InlineFragmentNode,
+    print_ast,
+)
 from ...type import is_composite_type
 from ...utilities import type_from_ast
 from . import ValidationRule
@@ -15,7 +19,7 @@ class FragmentsOnCompositeTypesRule(ValidationRule):
     must also be a composite type.
     """
 
-    def enter_inline_fragment(self, node: InlineFragmentNode, *_args):
+    def enter_inline_fragment(self, node: InlineFragmentNode, *_args) -> None:
         type_condition = node.type_condition
         if type_condition:
             type_ = type_from_ast(self.context.schema, type_condition)
@@ -29,7 +33,7 @@ def enter_inline_fragment(self, node: InlineFragmentNode, *_args):
                     )
                 )
 
-    def enter_fragment_definition(self, node: FragmentDefinitionNode, *_args):
+    def enter_fragment_definition(self, node: FragmentDefinitionNode, *_args) -> None:
         type_condition = node.type_condition
         type_ = type_from_ast(self.context.schema, type_condition)
         if type_ and not is_composite_type(type_):

src/graphql/validation/rules/known_argument_names.py

@@ -1,7 +1,13 @@
 from typing import cast, Dict, List, Union
 
 from ...error import GraphQLError
-from ...language import ArgumentNode, DirectiveDefinitionNode, DirectiveNode, SKIP
+from ...language import (
+    ArgumentNode,
+    DirectiveDefinitionNode,
+    DirectiveNode,
+    SKIP,
+    VisitorAction,
+)
 from ...pyutils import did_you_mean, suggestion_list
 from ...type import specified_directives
 from . import ASTValidationRule, SDLValidationContext, ValidationContext
@@ -37,7 +43,7 @@ def __init__(self, context: Union[ValidationContext, SDLValidationContext]):
 
         self.directive_args = directive_args
 
-    def enter_directive(self, directive_node: DirectiveNode, *_args):
+    def enter_directive(self, directive_node: DirectiveNode, *_args) -> VisitorAction:
         directive_name = directive_node.name.value
         known_args = self.directive_args.get(directive_name)
         if directive_node.arguments and known_args is not None:
@@ -67,7 +73,7 @@ class KnownArgumentNamesRule(KnownArgumentNamesOnDirectivesRule):
     def __init__(self, context: ValidationContext):
         super().__init__(context)
 
-    def enter_argument(self, arg_node: ArgumentNode, *args):
+    def enter_argument(self, arg_node: ArgumentNode, *args) -> None:
         context = self.context
         arg_def = context.get_argument()
         field_def = context.get_field_def()

src/graphql/validation/rules/known_directives.py

@@ -41,7 +41,9 @@ def __init__(self, context: Union[ValidationContext, SDLValidationContext]):
                 ]
         self.locations_map = locations_map
 
-    def enter_directive(self, node: DirectiveNode, _key, _parent, _path, ancestors):
+    def enter_directive(
+        self, node: DirectiveNode, _key, _parent, _path, ancestors
+    ) -> None:
         name = node.name.value
         locations = self.locations_map.get(name)
         if locations:

src/graphql/validation/rules/known_fragment_names.py

@@ -12,7 +12,7 @@ class KnownFragmentNamesRule(ValidationRule):
     fragments defined in the same document.
     """
 
-    def enter_fragment_spread(self, node: FragmentSpreadNode, *_args):
+    def enter_fragment_spread(self, node: FragmentSpreadNode, *_args) -> None:
         fragment_name = node.name.value
         fragment = self.context.get_fragment(fragment_name)
         if not fragment: