Merge pull request #68 from evo-company/typings-hiku.graph

[typings] hiku.graph

hiku/directives/__init__.py

@@ -1,12 +1,11 @@
 from typing import (
     Optional,
     Union,
+    TYPE_CHECKING
 )
 
-from hiku.graph import (
-    Field,
-    Link,
-)
+if TYPE_CHECKING:
+    from hiku.graph import Field, Link
 
 
 class DirectiveBase:
@@ -24,7 +23,7 @@ def accept(self, visitor):
         return visitor.visit_deprecated_directive(self)
 
 
-def get_deprecated(field: Union[Field, Link]) -> Optional[Deprecated]:
+def get_deprecated(field: Union['Field', 'Link']) -> Optional[Deprecated]:
     """Get deprecated directive"""
     return next(
         (d for d in field.directives if isinstance(d, Deprecated)),

hiku/engine.py

@@ -2,11 +2,20 @@
 import warnings
 import dataclasses
 
-from typing import Any
+from typing import (
+    Any,
+    TypeVar,
+    Callable,
+    cast,
+)
 from functools import partial
 from itertools import chain, repeat
 from collections import defaultdict
 from collections.abc import Sequence, Mapping, Hashable
+from typing_extensions import (
+    ParamSpec,
+    Concatenate,
+)
 
 from . import query as hiku_query
 from .graph import Link, Maybe, One, Many, Nothing, Field
@@ -326,7 +335,7 @@ def link_result_to_ids(from_list, link_type, result):
 
 class Query(Workflow):
 
-    def __init__(self, queue, task_set, graph, query, ctx):
+    def __init__(self, queue, task_set, graph, query, ctx: 'Context'):
         self._queue = queue
         self._task_set = task_set
         self._graph = graph
@@ -437,12 +446,23 @@ def _schedule_link(self, node, graph_link, query_link, ids):
         return dep
 
 
-def pass_context(func):
-    func.__pass_context__ = True
-    return func
+R = TypeVar('R')
+P = ParamSpec('P')
 
 
-def _do_pass_context(func):
+def pass_context(
+    func: Callable[P, R]
+) -> Callable[Concatenate['Context', P], R]:
+    """Decorator to pass context to a function as a first argument.
+
+    Can be used on functions for ``Field`` and ``Link``.
+    Can not be used on functions with ``@define`` decorator
+    """
+    func.__pass_context__ = True  # type: ignore[attr-defined]
+    return cast(Callable[Concatenate['Context', P], R], func)
+
+
+def _do_pass_context(func: Callable) -> bool:
     return getattr(func, '__pass_context__', False)
 
 

hiku/federation/endpoint.py

@@ -111,7 +111,7 @@ def execute(
         self, graph: Graph, op: Operation, ctx: Optional[Dict]
     ) -> Dict:
         stripped_query = _process_query(graph, op.query)
-        result = self.engine.execute(graph, stripped_query, ctx)
+        result = self.engine.execute(graph, stripped_query, ctx or {})
         return self.postprocess_result(result, graph, op)
 
     def dispatch(self, data: Dict) -> Dict:

hiku/federation/sdl.py

@@ -31,7 +31,7 @@
     AnyMeta,
     FloatMeta,
     BooleanMeta,
-    Record,
+    RecordMeta,
 )
 
 
@@ -111,7 +111,7 @@ def _encode_default_value(value) -> Optional[ast.ValueNode]:
 
 
 class Exporter(GraphVisitor):
-    def export_record(self, type_name: str, obj: Record):
+    def export_record(self, type_name: str, obj: RecordMeta):
         def new_field(name: str, type_):
             return ast.FieldDefinitionNode(
                 name=_name(name),