refactor: Split complexity into building the tree and then the evaluation to allow testing the describe function

Author: Checho3388

src/graphql_complexity/evaluator/complexity.py

@@ -1,16 +1,29 @@
 from graphql import GraphQLSchema, TypeInfo, TypeInfoVisitor, parse, visit
 
+from . import nodes
 from .visitor import ComplexityVisitor
 from ..config import Config
 from ..estimators import ComplexityEstimator
 
 
 def get_complexity(query: str, schema: GraphQLSchema, estimator: ComplexityEstimator, config: Config = None) -> int:
+    """Calculate the complexity of a query using the provided estimator."""
+    tree = build_complexity_tree(query, schema, estimator, config)
+
+    return tree.evaluate()
+
+
+def build_complexity_tree(
+        query: str,
+        schema: GraphQLSchema,
+        estimator: ComplexityEstimator,
+        config: Config | None = None,
+) -> nodes.ComplexityNode:
     """Calculate the complexity of a query using the provided estimator."""
     ast = parse(query)
     type_info = TypeInfo(schema)
 
     visitor = ComplexityVisitor(estimator=estimator, type_info=type_info, config=config)
     visit(ast, TypeInfoVisitor(type_info, visitor))
 
-    return visitor.evaluate()
+    return visitor.complexity_tree

src/graphql_complexity/evaluator/nodes.py

@@ -21,57 +21,54 @@ class ComplexityNode:
     parent: 'ComplexityNode' = None
     children: list['ComplexityNode'] = dataclasses.field(default_factory=list)
 
-    def evaluate(self, *args, **kwargs) -> int:
+    def evaluate(self) -> int:
         raise NotImplementedError
 
-    def describe(self, depth=0):
+    def describe(self, depth=0) -> str:
+        """Return a friendly representation of the node and its children complexity."""
         return (
             f"{chr(9) * depth}{self.name} ({self.__class__.__name__}) = {self.evaluate()}" +
             f"{chr(10) if self.children else ''}" +
             '\n'.join(c.describe(depth=depth+1) for c in self.children)
         )
 
-    def add_child(self, node: 'ComplexityNode'):
+    def add_child(self, node: 'ComplexityNode') -> None:
+        """Add a child to the current node."""
         self.children.append(node)
         node.parent = self
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
 class RootNode(ComplexityNode):
-    def evaluate(self, *args, **kwargs) -> int:
-        return sum(
-            child.evaluate(*args, **kwargs) for child in self.children
-        )
+    def evaluate(self) -> int:
+        return sum(child.evaluate() for child in self.children)
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
-class FragmentNode(ComplexityNode):
+class FragmentSpreadNode(ComplexityNode):
+    fragments_definition: dict
 
-    def evaluate(self, *, fragments_definition: dict[str, "ComplexityNode"]):
-        fragment = fragments_definition.get(self.name)
+    def evaluate(self):
+        fragment = self.fragments_definition.get(self.name)
         if not fragment:
             return 0
-        return fragment.evaluate(fragments_definition=fragments_definition)
+        return fragment.evaluate()
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
 class Field(ComplexityNode):
     complexity: int
 
-    def evaluate(self, *args, **kwargs) -> int:
-        return self.complexity + sum(
-            child.evaluate(*args, **kwargs) for child in self.children
-        )
+    def evaluate(self) -> int:
+        return self.complexity + sum(child.evaluate() for child in self.children)
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
 class ListField(Field):
     count: int = 1
 
-    def evaluate(self, *args, **kwargs) -> int:
-        return self.complexity + self.count * sum(
-            child.evaluate(*args, **kwargs) for child in self.children
-        )
+    def evaluate(self) -> int:
+        return self.complexity + self.count * sum(child.evaluate() for child in self.children)
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
@@ -90,14 +87,14 @@ def wrap(cls, node: ComplexityNode):
         node.parent.add_child(wrapper)
         return wrapper
 
-    def evaluate(self, *args, **kwargs) -> int:
+    def evaluate(self) -> int:
         return 0
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
 class MetaField(ComplexityNode):
 
-    def evaluate(self, *args, **kwargs) -> int:
+    def evaluate(self) -> int:
         return 0
 
 

src/graphql_complexity/evaluator/visitor.py

@@ -40,16 +40,17 @@ def __init__(
         self.variables = variables or {}
         self.type_info = type_info
         self.fragments: dict[str, nodes.ComplexityNode] = {}
-        self.current_node = nodes.RootNode(name="root")
+        self.root = nodes.RootNode(name="root")
+        self.current_node = self.root
         self._previous_current_node = None
-        self._ignore_until_leave = None
         super().__init__()
 
-    def evaluate(self) -> int:
-        """Evaluate the complexity of the operations after visiting the document."""
-        return self.current_node.evaluate(
-            fragments_definition=self.fragments
-        )
+    @property
+    def complexity_tree(self) -> nodes.ComplexityNode:
+        """Return the complexity tree after visiting the document.
+        The tree is represented by a RootNode with the complexity of the operations
+        represented as Node children. Each node is evaluated returning the complexity."""
+        return self.root
 
     def enter_variable_definition(self, node, key, parent, path, ancestors):
         input_variable = self.variables.get(node.variable.name.value)
@@ -85,7 +86,12 @@ def leave_fragment_definition(self, node, *args, **kwargs):
 
     def enter_fragment_spread(self, node, *args, **kwargs):
         """Add a lazy fragment to the current complexity list."""
-        self.current_node.add_child(nodes.FragmentNode(name=node.name.value))
+        self.current_node.add_child(
+            nodes.FragmentSpreadNode(
+                name=node.name.value,
+                fragments_definition=self.fragments
+            )
+        )
 
 
 def should_include_field(node: DirectiveNode, variables: dict[str, Any]) -> bool:

tests/test_tree.py

@@ -0,0 +1,77 @@
+import pytest
+from graphql import build_schema
+
+from graphql_complexity import SimpleEstimator
+from graphql_complexity.evaluator.complexity import build_complexity_tree
+from graphql_complexity.evaluator.nodes import ComplexityNode
+from tests import ut_utils
+
+
+def _build_complexity_tree(query: str, estimator=None):
+    estimator = estimator or SimpleEstimator(1)
+    schema = build_schema(ut_utils.schema)
+    return build_complexity_tree(query, schema, estimator)
+
+
+def test_tree_describes_simple_query():
+    query = """
+        query {
+            version
+        }
+    """
+
+    tree = _build_complexity_tree(query)
+
+    assert tree.describe() == """root (RootNode) = 1
+\tversion (Field) = 1"""
+
+
+def test_tree_describes_skipped_fields():
+    query = """query Foo ($shouldSkip: Boolean = false) {
+        version @include(if: $shouldSkip)
+    }"""
+
+    tree = _build_complexity_tree(query)
+
+    assert tree.describe() == """root (RootNode) = 0
+\tversion (SkippedField) = 0"""
+
+
+def test_tree_describes_lists():
+    query = """query {
+        droid {
+            friends {
+                name
+            }
+        }
+    }"""
+
+    tree = _build_complexity_tree(query)
+
+    assert tree.describe() == """root (RootNode) = 3
+\tdroid (Field) = 3
+\t\tfriends (ListField) = 2
+\t\t\tname (Field) = 1"""
+
+
+def test_describes_fragment_spread():
+    query = """query {
+        ...fragmentName
+    }
+    fragment fragmentName on Droid {
+        id
+    }"""
+
+    tree = _build_complexity_tree(query)
+
+    assert tree.describe() == """root (RootNode) = 1
+\tfragmentName (FragmentSpreadNode) = 1"""
+
+
+def test_complexity_node_can_not_be_evaluated():
+    node = ComplexityNode(
+        name="root"
+    )
+
+    with pytest.raises(NotImplementedError):
+        node.evaluate()