Extract check for possible extensions into a separate rule

Cito · Cito · commit 218ff57b8bbf · 2019-03-07T22:54:23.000+01:00
Replicates graphql/graphql-js@e6a3f08
diff --git a/README.md b/README.md
@@ -13,7 +13,7 @@ a query language for APIs created by Facebook.
 [![Code Style](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/ambv/black)
 
 The current version 1.0.1 of GraphQL-core-next is up-to-date with GraphQL.js version
-14.0.2. All parts of the API are covered by an extensive test suite of currently 1683
+14.0.2. All parts of the API are covered by an extensive test suite of currently 1690
 unit tests.
 
 
diff --git a/graphql/utilities/extend_schema.py b/graphql/utilities/extend_schema.py
@@ -3,19 +3,13 @@
 from itertools import chain
 from typing import Any, Callable, Dict, List, Optional, Union, Tuple, cast
 
-from ..error import GraphQLError
 from ..language import (
     DirectiveDefinitionNode,
     DocumentNode,
-    EnumTypeExtensionNode,
-    InputObjectTypeExtensionNode,
-    InterfaceTypeExtensionNode,
-    ObjectTypeExtensionNode,
     OperationType,
     SchemaExtensionNode,
     SchemaDefinitionNode,
     TypeDefinitionNode,
-    UnionTypeExtensionNode,
     TypeExtensionNode,
 )
 from ..type import (
@@ -110,17 +104,7 @@ def extend_schema(
             type_name = def_.name.value
             type_definition_map[type_name] = def_
         elif isinstance(def_, TypeExtensionNode):
-            # Sanity check that this type extension exists within the schema's existing
-            # types.
             extended_type_name = def_.name.value
-            existing_type = schema.get_type(extended_type_name)
-            if not existing_type:
-                raise GraphQLError(
-                    f"Cannot extend type '{extended_type_name}'"
-                    " because it does not exist in the existing schema.",
-                    [def_],
-                )
-            check_extension_node(existing_type, def_)
             type_extensions_map[extended_type_name].append(def_)
         elif isinstance(def_, DirectiveDefinitionNode):
             directive_definitions.append(def_)
@@ -502,25 +486,3 @@ def resolve_type(type_name: str) -> GraphQLNamedType:
         ast_node=schema.ast_node,
         extension_ast_nodes=schema_extension_ast_nodes,
     )
-
-
-def check_extension_node(type_: GraphQLNamedType, node: TypeExtensionNode):
-    if isinstance(node, ObjectTypeExtensionNode):
-        if not is_object_type(type_):
-            raise GraphQLError(f"Cannot extend non-object type '{type_.name}'.", [node])
-    elif isinstance(node, InterfaceTypeExtensionNode):
-        if not is_interface_type(type_):
-            raise GraphQLError(
-                f"Cannot extend non-interface type '{type_.name}'.", [node]
-            )
-    elif isinstance(node, EnumTypeExtensionNode):
-        if not is_enum_type(type_):
-            raise GraphQLError(f"Cannot extend non-enum type '{type_.name}'.", [node])
-    elif isinstance(node, UnionTypeExtensionNode):
-        if not is_union_type(type_):
-            raise GraphQLError(f"Cannot extend non-union type '{type_.name}'.", [node])
-    elif isinstance(node, InputObjectTypeExtensionNode):
-        if not is_input_object_type(type_):
-            raise GraphQLError(
-                f"Cannot extend non-input object type '{type_.name}'.", [node]
-            )
diff --git a/graphql/validation/rules/possible_type_extensions.py b/graphql/validation/rules/possible_type_extensions.py
@@ -0,0 +1,127 @@
+import re
+from functools import partial
+from typing import Any, List
+
+from ...error import GraphQLError
+from ...language import TypeDefinitionNode, TypeExtensionNode
+from ...pyutils import quoted_or_list, suggestion_list
+from ...type import (
+    is_enum_type,
+    is_input_object_type,
+    is_interface_type,
+    is_object_type,
+    is_scalar_type,
+    is_union_type,
+)
+from . import SDLValidationContext, SDLValidationRule
+
+__all__ = [
+    "PossibleTypeExtensionsRule",
+    "extending_unknown_type_message",
+    "extending_different_type_kind_message",
+]
+
+
+def extending_unknown_type_message(type_name: str, suggested_types: List[str]) -> str:
+    message = f"Cannot extend type '{type_name}' because it is not defined."
+    if suggested_types:
+        message += f" Did you mean {quoted_or_list(suggested_types)}?"
+    return message
+
+
+def extending_different_type_kind_message(type_name: str, kind: str) -> str:
+    return f"Cannot extend non-{kind} type {type_name}"
+
+
+class PossibleTypeExtensionsRule(SDLValidationRule):
+    """Possible type extension
+
+     A type extension is only valid if the type is defined and has the same kind.
+    """
+
+    def __init__(self, context: SDLValidationContext) -> None:
+        super().__init__(context)
+        self.schema = context.schema
+        self.defined_types = {
+            def_.name.value: def_
+            for def_ in context.document.definitions
+            if isinstance(def_, TypeDefinitionNode)
+        }
+
+    def check_extension(self, node: TypeExtensionNode, *_args):
+        schema = self.schema
+        type_name = node.name.value
+        def_node = self.defined_types.get(type_name)
+        existing_type = schema.get_type(type_name) if schema else None
+
+        if def_node:
+            expected_kind = def_kind_to_ext_kind(def_node.kind)
+            if expected_kind != node.kind:
+                self.report_error(
+                    GraphQLError(
+                        extending_different_type_kind_message(
+                            type_name, extension_kind_to_type_name(expected_kind)
+                        ),
+                        [def_node, node],
+                    )
+                )
+        elif existing_type:
+            expected_kind = type_to_ext_kind(existing_type)
+            if expected_kind != node.kind:
+                self.report_error(
+                    GraphQLError(
+                        extending_different_type_kind_message(
+                            type_name, extension_kind_to_type_name(expected_kind)
+                        ),
+                        [node],
+                    )
+                )
+        else:
+            all_type_names = list(self.defined_types)
+            if self.schema:
+                all_type_names.extend(self.schema.type_map)
+            suggested_types = suggestion_list(type_name, all_type_names)
+            self.report_error(
+                GraphQLError(
+                    extending_unknown_type_message(type_name, suggested_types),
+                    [node.name],
+                )
+            )
+
+    enter_scalar_type_extension = enter_object_type_extension = check_extension
+    enter_interface_type_extension = enter_union_type_extension = check_extension
+    enter_enum_type_extension = enter_input_object_type_extension = check_extension
+
+
+def_kind_to_ext_kind = partial(re.compile("(?<=_type_)definition$").sub, "extension")
+
+
+def type_to_ext_kind(type_: Any) -> str:
+    if is_scalar_type(type_):
+        return "scalar_type_extension"
+    elif is_object_type(type_):
+        return "object_type_extension"
+    elif is_interface_type(type_):
+        return "interface_type_extension"
+    elif is_union_type(type_):
+        return "union_type_extension"
+    elif is_enum_type(type_):
+        return "enum_type_extension"
+    elif is_input_object_type(type_):
+        return "input_object_type_extension"
+    else:
+        return "unknown_type_extension"
+
+
+_type_names_for_extension_kinds = {
+    "scalar_type_extension": "scalar",
+    "object_type_extension": "object",
+    "interface_type_extension": "interface",
+    "union_type_extension": "union",
+    "enum_type_extension": "enum",
+    "input_object_type_extension": "input object",
+}
+
+
+def extension_kind_to_type_name(kind: str) -> str:
+    return _type_names_for_extension_kinds.get(kind, "unknown type")
diff --git a/graphql/validation/specified_rules.py b/graphql/validation/specified_rules.py
@@ -87,6 +87,7 @@
 from .rules.unique_enum_value_names import UniqueEnumValueNamesRule
 from .rules.unique_field_definition_names import UniqueFieldDefinitionNamesRule
 from .rules.unique_directive_names import UniqueDirectiveNamesRule
+from .rules.possible_type_extensions import PossibleTypeExtensionsRule
 from .rules.known_argument_names import KnownArgumentNamesOnDirectivesRule
 from .rules.provided_required_arguments import ProvidedRequiredArgumentsOnDirectivesRule
 
@@ -137,6 +138,7 @@
     KnownTypeNamesRule,
     KnownDirectivesRule,
     UniqueDirectivesPerLocationRule,
+    PossibleTypeExtensionsRule,
     KnownArgumentNamesOnDirectivesRule,
     UniqueArgumentNamesRule,
     UniqueInputFieldNamesRule,
diff --git a/tests/utilities/test_extend_schema.py b/tests/utilities/test_extend_schema.py
@@ -1,7 +1,6 @@
 from pytest import raises
 
 from graphql import graphql_sync
-from graphql.error import GraphQLError
 from graphql.language import parse, print_ast, DirectiveLocation, DocumentNode
 from graphql.pyutils import dedent
 from graphql.type import (
@@ -1068,62 +1067,6 @@ def does_not_allow_replacing_an_existing_enum_value():
             " It cannot also be defined in this type extension."
         )
 
-    def does_not_allow_extending_an_unknown_type():
-        for sdl in [
-            "extend scalar UnknownType @foo",
-            "extend type UnknownType @foo",
-            "extend interface UnknownType @foo",
-            "extend enum UnknownType @foo",
-            "extend union UnknownType @foo",
-            "extend input UnknownType @foo",
-        ]:
-            with raises(GraphQLError) as exc_info:
-                extend_test_schema(sdl)
-            assert str(exc_info.value).startswith(
-                "Cannot extend type 'UnknownType'"
-                " because it does not exist in the existing schema."
-            )
-
-    def it_does_not_allow_extending_a_mismatch_type():
-        type_sdl = """
-            extend type SomeInterface @foo
-            """
-        with raises(GraphQLError) as exc_info:
-            extend_test_schema(type_sdl)
-        assert str(exc_info.value).startswith(
-            "Cannot extend non-object type 'SomeInterface'."
-        )
-
-        interface_sdl = """
-            extend interface Foo @foo
-        """
-        with raises(GraphQLError) as exc_info:
-            extend_test_schema(interface_sdl)
-        assert str(exc_info.value).startswith("Cannot extend non-interface type 'Foo'.")
-
-        enum_sdl = """
-            extend enum Foo @foo
-        """
-        with raises(GraphQLError) as exc_info:
-            extend_test_schema(enum_sdl)
-        assert str(exc_info.value).startswith("Cannot extend non-enum type 'Foo'.")
-
-        union_sdl = """
-            extend union Foo @foo
-        """
-        with raises(GraphQLError) as exc_info:
-            extend_test_schema(union_sdl)
-        assert str(exc_info.value).startswith("Cannot extend non-union type 'Foo'.")
-
-        input_sdl = """
-            extend input Foo @foo
-        """
-        with raises(GraphQLError) as exc_info:
-            extend_test_schema(input_sdl)
-        assert str(exc_info.value).startswith(
-            "Cannot extend non-input object type 'Foo'."
-        )
-
     def describe_can_add_additional_root_operation_types():
         def does_not_automatically_include_common_root_type_names():
             schema = extend_test_schema(
diff --git a/tests/validation/test_possible_type_extensions.py b/tests/validation/test_possible_type_extensions.py
