refactor: Big refactor to support include and skip directives

Author: Checho3388

src/graphql_complexity/__init__.py

@@ -0,0 +1,13 @@
+from .estimators import (
+    ComplexityEstimator,
+    DirectivesEstimator,
+    SimpleEstimator
+)
+from .visitor import ComplexityVisitor
+
+__all__ = [
+    "ComplexityVisitor",
+    "SimpleEstimator",
+    "ComplexityEstimator",
+    "DirectivesEstimator",
+]

src/graphql_complexity/estimators.py

@@ -0,0 +1,5 @@
+from .base import ComplexityEstimator
+from .directive import DirectivesEstimator
+from .simple import SimpleEstimator
+
+__all__ = ["ComplexityEstimator", "SimpleEstimator", "DirectivesEstimator"]

src/graphql_complexity/estimators/__init__.py

@@ -0,0 +1,11 @@
+import abc
+
+
+class ComplexityEstimator(abc.ABC):
+    @abc.abstractmethod
+    def get_field_complexity(self, node, key, parent, path, ancestors) -> int:
+        """Return the complexity of the field."""
+
+    @abc.abstractmethod
+    def get_field_multiplier(self, node, key, parent, path, ancestors) -> int:
+        """Return the multiplier that will be applied to the children of the given node."""

src/graphql_complexity/estimators/base.py

@@ -0,0 +1,77 @@
+from graphql import Visitor, parse, visit
+
+from graphql_complexity.estimators.base import ComplexityEstimator
+
+DIRECTIVE_ESTIMATOR_FIELD_COMPLEXITY_NAME = "value"
+DEFAULT_COMPLEXITY_DIRECTIVE_NAME = "complexity"
+DEFAULT_COMPLEXITY_VALUE = 1
+
+
+class DirectivesVisitor(Visitor):
+    def __init__(self, directive_name: str, collector: dict[str, int]):
+        self._collector = collector
+        self.directive_name = directive_name
+        super().__init__()
+
+    def enter_field_definition(self, node, key, parent, path, ancestors):
+        for directive in node.directives:
+            if directive.name.value == self.directive_name:
+                complexity = next(
+                    arg.value.value
+                    for arg in directive.arguments
+                    if arg.name.value == DIRECTIVE_ESTIMATOR_FIELD_COMPLEXITY_NAME
+                )
+                self._collector[node.name.value] = int(complexity)
+                break
+
+
+class DirectivesEstimator(ComplexityEstimator):
+    """Complexity estimator that uses directives to get the complexity of the fields.
+    The complexity is calculated by visiting the schema and using the complexity
+    directive to get the complexity of each field.
+
+    Example:
+        Given the following schema:
+        ```qgl
+            directive @complexity(
+              value: Int!
+            ) on FIELD_DEFINITION
+
+            type Query {
+              oneField: String @complexity(value: 5)
+              otherField: String @complexity(value: 1)
+            }
+        ```
+        The complexity of the fields will be:
+            - oneField: 5
+            - otherField: 1
+        And the total complexity will be 6.
+    """
+
+    def __init__(
+        self,
+        schema: str,
+        directive_name: str = DEFAULT_COMPLEXITY_DIRECTIVE_NAME,
+        missing_complexity: int = DEFAULT_COMPLEXITY_VALUE,
+    ):
+        self.__directive_name = directive_name
+        self.__missing_complexity = int(missing_complexity)
+        self.__complexity_map = self.collect_from_schema(
+            schema=schema, directive_name=directive_name
+        )
+        super().__init__()
+
+    @staticmethod
+    def collect_from_schema(schema: str, directive_name: str) -> dict[str, int]:
+        collector = {}
+        ast = parse(schema)
+        visitor = DirectivesVisitor(collector=collector, directive_name=directive_name)
+        visit(ast, visitor)
+        return collector
+
+    def get_field_complexity(self, node, key, parent, path, ancestors) -> int:
+        return self.__complexity_map.get(node.name.value, self.__missing_complexity)
+
+    def get_field_multiplier(self, node, key, parent, path, ancestors) -> int:
+        # ToDo: Implement this method
+        return 1

src/graphql_complexity/estimators/directive.py

@@ -0,0 +1,42 @@
+from graphql_complexity.estimators.base import ComplexityEstimator
+
+
+class SimpleEstimator(ComplexityEstimator):
+    """Simple complexity estimator that returns a constant complexity and multiplier for all fields.
+    Constants can be set in the constructor.
+
+    Example:
+        Given the following query:
+        ```qgl
+            query {
+                user {
+                    name
+                    email
+                }
+            }
+        ```
+        As the complexity and multiplier are constant, the complexity of the fields will be:
+            - user: 1 * 1 = 1
+            - name: 1 * 1 = 1
+            - email: 1 * 1 = 1
+        And the total complexity will be 3.
+    """
+
+    def __init__(self, complexity: int = 1, multiplier: int = 1):
+        if complexity < 0:
+            raise ValueError(
+                "'complexity' must be a positive integer (greater or equal than 0)"
+            )
+        if multiplier < 0:
+            raise ValueError(
+                "'multiplier' must be a positive integer (greater or equal than 0)"
+            )
+        self.__complexity_constant = complexity
+        self.__multiplier_constant = multiplier
+        super().__init__()
+
+    def get_field_complexity(self, node, key, parent, path, ancestors) -> int:
+        return self.__complexity_constant
+
+    def get_field_multiplier(self, node, key, parent, path, ancestors) -> int:
+        return self.__multiplier_constant

src/graphql_complexity/estimators/simple.py

@@ -0,0 +1,41 @@
+from typing import Type
+
+from graphql import GraphQLError, visit
+from strawberry.extensions import SchemaExtension
+
+from graphql_complexity import (
+    ComplexityEstimator,
+    ComplexityVisitor,
+    SimpleEstimator
+)
+
+
+def build_complexity_extension(
+    estimator: ComplexityEstimator | None = None,
+    max_complexity: int | None = None,
+) -> Type[SchemaExtension]:
+    estimator = estimator or SimpleEstimator(1, 1)
+
+    class ComplexityExtension(SchemaExtension):
+        visitor = None
+        estimated_complexity: int = None
+
+        def on_validate(
+            self,
+        ):
+            self.visitor = ComplexityVisitor(estimator=estimator)
+            visit(self.execution_context.graphql_document, self.visitor)
+
+            self.estimated_complexity = self.visitor.evaluate()
+
+            if max_complexity and self.estimated_complexity > max_complexity:
+                error = GraphQLError(
+                    f"Query is too complex. Max complexity is {max_complexity}, estimated "
+                    f"complexity is {self.estimated_complexity}"
+                )
+                self.execution_context.errors = [error]
+
+        def get_results(self):
+            return {"complexity": {"value": self.estimated_complexity}}
+
+    return ComplexityExtension

src/graphql_complexity/extensions.py

@@ -0,0 +1,10 @@
+from graphql import parse, visit
+
+from . import ComplexityEstimator, ComplexityVisitor
+
+
+def get_complexity(query: str, estimator: ComplexityEstimator) -> int:
+    ast = parse(query)
+    visitor = ComplexityVisitor(estimator=estimator)
+    visit(ast, visitor)
+    return visitor.evaluate()