Use VisitorMeta to power TypeInfo so that it's not a huge isinstance chain.

jhgg · jhgg · commit 310d0a91a830 · 2015-10-07T15:48:07.000-04:00
diff --git a/graphql/core/language/visitor.py b/graphql/core/language/visitor.py
@@ -1,50 +1,9 @@
 from copy import copy
+
 import six
 from . import ast
+from .visitor_meta import VisitorMeta, QUERY_DOCUMENT_KEYS
 
-QUERY_DOCUMENT_KEYS = {
-    ast.Name: (),
-
-    ast.Document: ('definitions',),
-    ast.OperationDefinition: ('name', 'variable_definitions', 'directives', 'selection_set'),
-    ast.VariableDefinition: ('variable', 'type', 'default_value'),
-    ast.Variable: ('name',),
-    ast.SelectionSet: ('selections',),
-    ast.Field: ('alias', 'name', 'arguments', 'directives', 'selection_set'),
-    ast.Argument: ('name', 'value'),
-
-    ast.FragmentSpread: ('name', 'directives'),
-    ast.InlineFragment: ('type_condition', 'directives', 'selection_set'),
-    ast.FragmentDefinition: ('name', 'type_condition', 'directives', 'selection_set'),
-
-    ast.IntValue: (),
-    ast.FloatValue: (),
-    ast.StringValue: (),
-    ast.BooleanValue: (),
-    ast.EnumValue: (),
-    ast.ListValue: ('values',),
-    ast.ObjectValue: ('fields',),
-    ast.ObjectField: ('name', 'value'),
-
-    ast.Directive: ('name', 'arguments'),
-
-    ast.NamedType: ('name',),
-    ast.ListType: ('type',),
-    ast.NonNullType: ('type',),
-
-    ast.ObjectTypeDefinition: ('name', 'interfaces', 'fields'),
-    ast.FieldDefinition: ('name', 'arguments', 'type'),
-    ast.InputValueDefinition: ('name', 'type', 'defaultValue'),
-    ast.InterfaceTypeDefinition: ('name', 'fields'),
-    ast.UnionTypeDefinition: ('name', 'types'),
-    ast.ScalarTypeDefinition: ('name',),
-    ast.EnumTypeDefinition: ('name', 'values'),
-    ast.EnumValueDefinition: ('name',),
-    ast.InputObjectTypeDefinition: ('name', 'fields'),
-    ast.TypeExtensionDefinition: ('definition',),
-}
-
-AST_KIND_TO_TYPE = {c.__name__: c for c in QUERY_DOCUMENT_KEYS.keys()}
 
 BREAK = object()
 REMOVE = object()
@@ -189,33 +148,6 @@ def visit(root, visitor, key_map=None):
     return new_root
 
 
-class VisitorMeta(type):
-    def __new__(cls, name, bases, attrs):
-        enter_handlers = {}
-        leave_handlers = {}
-
-        for base in bases:
-            if hasattr(base, '_enter_handlers'):
-                enter_handlers.update(base._enter_handlers)
-            if hasattr(base, '_leave_handlers'):
-                leave_handlers.update(base._leave_handlers)
-
-        for attr, val in attrs.items():
-            if attr.startswith('enter_'):
-                ast_kind = attr[6:]
-                ast_type = AST_KIND_TO_TYPE.get(ast_kind)
-                enter_handlers[ast_type] = val
-
-            elif attr.startswith('leave_'):
-                ast_kind = attr[6:]
-                ast_type = AST_KIND_TO_TYPE.get(ast_kind)
-                leave_handlers[ast_type] = val
-
-        attrs['_get_enter_handler'] = enter_handlers.get
-        attrs['_get_leave_handler'] = leave_handlers.get
-        return super(VisitorMeta, cls).__new__(cls, name, bases, attrs)
-
-
 @six.add_metaclass(VisitorMeta)
 class Visitor(object):
     def enter(self, node, key, parent, path, ancestors):
diff --git a/graphql/core/language/visitor_meta.py b/graphql/core/language/visitor_meta.py
@@ -0,0 +1,72 @@
+from . import ast
+
+QUERY_DOCUMENT_KEYS = {
+    ast.Name: (),
+
+    ast.Document: ('definitions',),
+    ast.OperationDefinition: ('name', 'variable_definitions', 'directives', 'selection_set'),
+    ast.VariableDefinition: ('variable', 'type', 'default_value'),
+    ast.Variable: ('name',),
+    ast.SelectionSet: ('selections',),
+    ast.Field: ('alias', 'name', 'arguments', 'directives', 'selection_set'),
+    ast.Argument: ('name', 'value'),
+
+    ast.FragmentSpread: ('name', 'directives'),
+    ast.InlineFragment: ('type_condition', 'directives', 'selection_set'),
+    ast.FragmentDefinition: ('name', 'type_condition', 'directives', 'selection_set'),
+
+    ast.IntValue: (),
+    ast.FloatValue: (),
+    ast.StringValue: (),
+    ast.BooleanValue: (),
+    ast.EnumValue: (),
+    ast.ListValue: ('values',),
+    ast.ObjectValue: ('fields',),
+    ast.ObjectField: ('name', 'value'),
+
+    ast.Directive: ('name', 'arguments'),
+
+    ast.NamedType: ('name',),
+    ast.ListType: ('type',),
+    ast.NonNullType: ('type',),
+
+    ast.ObjectTypeDefinition: ('name', 'interfaces', 'fields'),
+    ast.FieldDefinition: ('name', 'arguments', 'type'),
+    ast.InputValueDefinition: ('name', 'type', 'defaultValue'),
+    ast.InterfaceTypeDefinition: ('name', 'fields'),
+    ast.UnionTypeDefinition: ('name', 'types'),
+    ast.ScalarTypeDefinition: ('name',),
+    ast.EnumTypeDefinition: ('name', 'values'),
+    ast.EnumValueDefinition: ('name',),
+    ast.InputObjectTypeDefinition: ('name', 'fields'),
+    ast.TypeExtensionDefinition: ('definition',),
+}
+
+AST_KIND_TO_TYPE = {c.__name__: c for c in QUERY_DOCUMENT_KEYS.keys()}
+
+
+class VisitorMeta(type):
+    def __new__(cls, name, bases, attrs):
+        enter_handlers = {}
+        leave_handlers = {}
+
+        for base in bases:
+            if hasattr(base, '_enter_handlers'):
+                enter_handlers.update(base._enter_handlers)
+            if hasattr(base, '_leave_handlers'):
+                leave_handlers.update(base._leave_handlers)
+
+        for attr, val in attrs.items():
+            if attr.startswith('enter_'):
+                ast_kind = attr[6:]
+                ast_type = AST_KIND_TO_TYPE.get(ast_kind)
+                enter_handlers[ast_type] = val
+
+            elif attr.startswith('leave_'):
+                ast_kind = attr[6:]
+                ast_type = AST_KIND_TO_TYPE.get(ast_kind)
+                leave_handlers[ast_type] = val
+
+        attrs['_get_enter_handler'] = enter_handlers.get
+        attrs['_get_leave_handler'] = leave_handlers.get
+        return super(VisitorMeta, cls).__new__(cls, name, bases, attrs)
diff --git a/graphql/core/utils/type_info.py b/graphql/core/utils/type_info.py
@@ -1,4 +1,6 @@
-from ..language import ast
+import six
+
+from ..language import ast, visitor_meta
 from ..type.definition import (
     GraphQLInputObjectType,
     GraphQLList,
@@ -16,6 +18,8 @@ def pop(lst):
         lst.pop()
 
 
+# noinspection PyPep8Naming
+@six.add_metaclass(visitor_meta.VisitorMeta)
 class TypeInfo(object):
     def __init__(self, schema):
         self._schema = schema
@@ -48,80 +52,105 @@ def get_directive(self):
     def get_argument(self):
         return self._argument
 
+    def leave(self, node):
+        method = self._get_leave_handler(type(node))
+        if method:
+            return method(self)
+
     def enter(self, node):
-        schema = self._schema
-        type = None
-        if isinstance(node, ast.SelectionSet):
-            named_type = get_named_type(self.get_type())
-            composite_type = None
-            if is_composite_type(named_type):
-                composite_type = named_type
-            self._parent_type_stack.append(composite_type)
-        elif isinstance(node, ast.Field):
-            parent_type = self.get_parent_type()
-            field_def = None
-            if parent_type:
-                field_def = get_field_def(schema, parent_type, node)
-            self._field_def_stack.append(field_def)
-            self._type_stack.append(field_def and field_def.type)
-        elif isinstance(node, ast.Directive):
-            self._directive = schema.get_directive(node.name.value)
-        elif isinstance(node, ast.OperationDefinition):
-            if node.operation == 'query':
-                type = schema.get_query_type()
-            elif node.operation == 'mutation':
-                type = schema.get_mutation_type()
-            self._type_stack.append(type)
-        elif isinstance(node, (ast.InlineFragment, ast.FragmentDefinition)):
-            type_condition_ast = node.type_condition
-            type = type_from_ast(schema, type_condition_ast) if type_condition_ast else self.get_type()
-            self._type_stack.append(type)
-        elif isinstance(node, ast.VariableDefinition):
-            self._input_type_stack.append(type_from_ast(schema, node.type))
-        elif isinstance(node, ast.Argument):
-            arg_def = None
-            arg_type = None
-            field_or_directive = self.get_directive() or self.get_field_def()
-            if field_or_directive:
-                arg_def = [arg for arg in field_or_directive.args if arg.name == node.name.value]
-                if arg_def:
-                    arg_def = arg_def[0]
-                    arg_type = arg_def.type
-                else:
-                    arg_def = None
-            self._argument = arg_def
-            self._input_type_stack.append(arg_type)
-        elif isinstance(node, ast.ListValue):
-            list_type = get_nullable_type(self.get_input_type())
-            self._input_type_stack.append(
-                list_type.of_type if isinstance(list_type, GraphQLList) else None
-            )
-        elif isinstance(node, ast.ObjectField):
-            object_type = get_named_type(self.get_input_type())
-            field_type = None
-            if isinstance(object_type, GraphQLInputObjectType):
-                input_field = object_type.get_fields().get(node.name.value)
-                field_type = input_field.type if input_field else None
-            self._input_type_stack.append(field_type)
+        method = self._get_enter_handler(type(node))
+        if method:
+            return method(self, node)
 
-    def leave(self, node):
-        if isinstance(node, ast.SelectionSet):
-            pop(self._parent_type_stack)
-        elif isinstance(node, ast.Field):
-            pop(self._field_def_stack)
-            pop(self._type_stack)
-        elif isinstance(node, ast.Directive):
-            self._directive = None
-        elif isinstance(node, (
-                ast.OperationDefinition,
-                ast.InlineFragment,
-                ast.FragmentDefinition,
-        )):
-            pop(self._type_stack)
-        elif isinstance(node, ast.VariableDefinition):
-            pop(self._input_type_stack)
-        elif isinstance(node, ast.Argument):
-            self._argument = None
-            pop(self._input_type_stack)
-        elif isinstance(node, (ast.ListType, ast.ObjectField)):
-            pop(self._input_type_stack)
+    def enter_SelectionSet(self, node):
+        named_type = get_named_type(self.get_type())
+        composite_type = None
+        if is_composite_type(named_type):
+            composite_type = named_type
+        self._parent_type_stack.append(composite_type)
+
+    def enter_Field(self, node):
+        parent_type = self.get_parent_type()
+        field_def = None
+        if parent_type:
+            field_def = get_field_def(self._schema, parent_type, node)
+        self._field_def_stack.append(field_def)
+        self._type_stack.append(field_def and field_def.type)
+
+    def enter_Directive(self, node):
+        self._directive = self._schema.get_directive(node.name.value)
+
+    def enter_OperationDefinition(self, node):
+        definition_type = None
+        if node.operation == 'query':
+            definition_type = self._schema.get_query_type()
+        elif node.operation == 'mutation':
+            definition_type = self._schema.get_mutation_type()
+
+        self._type_stack.append(definition_type)
+
+    def enter_InlineFragment(self, node):
+        type_condition_ast = node.type_condition
+        type = type_from_ast(self._schema, type_condition_ast) if type_condition_ast else self.get_type()
+        self._type_stack.append(type)
+
+    enter_FragmentDefinition = enter_InlineFragment
+
+    def enter_VariableDefinition(self, node):
+        self._input_type_stack.append(type_from_ast(self._schema, node.type))
+
+    def enter_Argument(self, node):
+        arg_def = None
+        arg_type = None
+        field_or_directive = self.get_directive() or self.get_field_def()
+        if field_or_directive:
+            arg_def = [arg for arg in field_or_directive.args if arg.name == node.name.value]
+            if arg_def:
+                arg_def = arg_def[0]
+                arg_type = arg_def.type
+            else:
+                arg_def = None
+        self._argument = arg_def
+        self._input_type_stack.append(arg_type)
+
+    def enter_ListValue(self, node):
+        list_type = get_nullable_type(self.get_input_type())
+        self._input_type_stack.append(
+            list_type.of_type if isinstance(list_type, GraphQLList) else None
+        )
+
+    def enter_ObjectField(self, node):
+        object_type = get_named_type(self.get_input_type())
+        field_type = None
+        if isinstance(object_type, GraphQLInputObjectType):
+            input_field = object_type.get_fields().get(node.name.value)
+            field_type = input_field.type if input_field else None
+        self._input_type_stack.append(field_type)
+
+    def leave_SelectionSet(self):
+        pop(self._parent_type_stack)
+
+    def leave_Field(self):
+        pop(self._field_def_stack)
+        pop(self._type_stack)
+
+    def leave_Directive(self):
+        self._directive = None
+
+    def leave_OperationDefinition(self):
+        pop(self._type_stack)
+
+    leave_InlineFragment = leave_OperationDefinition
+    leave_FragmentDefinition = leave_OperationDefinition
+
+    def leave_VariableDefinition(self):
+        pop(self._input_type_stack)
+
+    def leave_Argument(self):
+        self._argument = None
+        pop(self._input_type_stack)
+
+    def leave_ListType(self):
+        pop(self._input_type_stack)
+
+    leave_ObjectField = leave_ListType
