add types

Author: suejung-sentry

graphql_api/tests/test_validation.py

@@ -23,7 +23,6 @@ def validate_query(query, *rules):
     return validate(schema, ast, rules=rules)
 
 
-# Tests for MaxDepthRule
 def test_max_depth_rule_allows_within_depth():
     query = """
     query {

graphql_api/validation.py

@@ -1,20 +1,25 @@
+from typing import Any, Type
+
 from graphql import GraphQLError, ValidationRule
 from graphql.language.ast import DocumentNode, FieldNode, OperationDefinitionNode
+from graphql.validation import ValidationContext
 
 
-def create_max_depth_rule(max_depth: int):
+def create_max_depth_rule(max_depth: int) -> Type[ValidationRule]:
     class MaxDepthRule(ValidationRule):
-        def __init__(self, context):
+        def __init__(self, context: ValidationContext) -> None:
             super().__init__(context)
-            self.operation_depth = 1
-            self.max_depth_reached = False
-            self.max_depth = max_depth
+            self.operation_depth: int = 1
+            self.max_depth_reached: bool = False
+            self.max_depth: int = max_depth
 
-        def enter_operation_definition(self, node: OperationDefinitionNode, *_args):
+        def enter_operation_definition(
+            self, node: OperationDefinitionNode, *_args: Any
+        ) -> None:
             self.operation_depth = 1
             self.max_depth_reached = False
 
-        def enter_field(self, node: FieldNode, *_args):
+        def enter_field(self, node: FieldNode, *_args: Any) -> None:
             self.operation_depth += 1
 
             if self.operation_depth > self.max_depth and not self.max_depth_reached:
@@ -26,25 +31,25 @@ def enter_field(self, node: FieldNode, *_args):
                     )
                 )
 
-        def leave_field(self, node: FieldNode, *_args):
+        def leave_field(self, node: FieldNode, *_args: Any) -> None:
             self.operation_depth -= 1
 
     return MaxDepthRule
 
 
-def create_max_aliases_rule(max_aliases: int):
+def create_max_aliases_rule(max_aliases: int) -> Type[ValidationRule]:
     class MaxAliasesRule(ValidationRule):
-        def __init__(self, context):
+        def __init__(self, context: ValidationContext) -> None:
             super().__init__(context)
-            self.alias_count = 0
-            self.has_reported_error = False
-            self.max_aliases = max_aliases
+            self.alias_count: int = 0
+            self.has_reported_error: bool = False
+            self.max_aliases: int = max_aliases
 
-        def enter_document(self, node: DocumentNode, *_args):
+        def enter_document(self, node: DocumentNode, *_args: Any) -> None:
             self.alias_count = 0
             self.has_reported_error = False
 
-        def enter_field(self, node: FieldNode, *_args):
+        def enter_field(self, node: FieldNode, *_args: Any) -> None:
             if node.alias:
                 self.alias_count += 1