Merge pull request #4 from seandstewart/seandstewart/ref-detection

fix: remove use of `graphlib.TypeNode` in type context

Author: seandstewart

src/typelib/ctx.py

@@ -2,35 +2,29 @@
 
 from __future__ import annotations
 
-import dataclasses
+import typing as tp
+import typing_extensions as te
 
-from typelib import graph
-from typelib.py import inspection, refs
+from typelib.py import refs
 
+ValueT = tp.TypeVar("ValueT")
+KeyT = te.TypeAliasType("KeyT", "type | refs.ForwardRef")
 
-class TypeContext(dict):
+
+class TypeContext(dict[KeyT, ValueT], tp.Generic[ValueT]):
     """A key-value mapping which can map between forward references and real types."""
 
-    def __missing__(self, key: graph.TypeNode | type | refs.ForwardRef):
+    def __missing__(self, key: type | refs.ForwardRef):
         """Hook to handle missing type references.
 
         Allows for sharing lookup results between forward references and real types.
 
         Args:
             key: The type or reference.
         """
-        # Eager wrap in a TypeNode
-        if not isinstance(key, graph.TypeNode):
-            key = graph.TypeNode(type=key)
-            return self[key]
-
         # If we missed a ForwardRef, we've already tried this, bail out.
-        type = key.type
-        if isinstance(type, refs.ForwardRef):
+        if isinstance(key, refs.ForwardRef):
             raise KeyError(key)
 
-        ref = refs.forwardref(
-            inspection.qualname(type), module=getattr(type, "__module__", None)
-        )
-        node = dataclasses.replace(key, type=ref)
-        return self[node]
+        ref = refs.forwardref(key)
+        return self[ref]

src/typelib/marshals/api.py

@@ -46,26 +46,24 @@ def marshaller(
             [`typing.ForwardRef`][], or string reference.
     """
     nodes = graph.static_order(t)
-    context: dict[type | graph.TypeNode, routines.AbstractMarshaller] = (
-        ctx.TypeContext()
-    )
+    context: ctx.TypeContext[routines.AbstractMarshaller] = ctx.TypeContext()
     if not nodes:
         return routines.NoOpMarshaller(t=t, context=context, var=None)  # type: ignore[arg-type]
 
     # "root" type will always be the final node in the sequence.
     root = nodes[-1]
     for node in nodes:
-        context[node] = _get_unmarshaller(node, context=context)
+        context[node.type] = _get_unmarshaller(node, context=context)
 
-    return context[root]
+    return context[root.type]
 
 
 def _get_unmarshaller(  # type: ignore[return]
     node: graph.TypeNode,
     context: routines.ContextT,
 ) -> routines.AbstractMarshaller[T]:
-    if node in context:
-        return context[node]
+    if node.type in context:
+        return context[node.type]
 
     for check, unmarshaller_cls in _HANDLERS.items():
         if check(node.type):

src/typelib/marshals/routines.py

@@ -14,7 +14,7 @@
 import uuid
 import warnings
 
-from typelib import graph, serdes
+from typelib import ctx, serdes
 from typelib.py import compat, inspection, refs
 
 T = tp.TypeVar("T")
@@ -93,7 +93,7 @@ def __init__(self, t: type[T], context: ContextT, *, var: str | None = None):
     def __call__(self, val: T) -> serdes.MarshalledValueT: ...
 
 
-ContextT: tp.TypeAlias = "tp.Mapping[type | graph.TypeNode, AbstractMarshaller]"
+ContextT: tp.TypeAlias = "ctx.TypeContext[AbstractMarshaller]"
 
 
 class NoOpMarshaller(AbstractMarshaller[T], tp.Generic[T]):
@@ -463,25 +463,21 @@ def __init__(self, t: type[_ST], context: ContextT, *, var: str | None = None):
 
     def _fields_by_var(self):
         fields_by_var = {}
-        tp_var_map = {(t.type, t.var): m for t, m in self.context.items()}
         hints = inspection.cached_type_hints(self.t)
         for name, hint in hints.items():
             resolved = refs.evaluate(hint)
-            fkey = (hint, name)
-            rkey = (resolved, name)
-            if fkey in tp_var_map:
-                fields_by_var[name] = tp_var_map[fkey]
-                continue
-            if rkey in tp_var_map:
-                fields_by_var[name] = tp_var_map[rkey]
+            m = self.context.get(hint) or self.context.get(resolved)
+            if m is None:
+                warnings.warn(
+                    "Failed to identify an unmarshaller for the associated type-variable pair: "
+                    f"Original ref: {hint}, Resolved ref: {resolved}. Will default to no-op.",
+                    stacklevel=4,
+                )
+                fields_by_var[name] = NoOpMarshaller(hint, self.context, var=name)
                 continue
 
-            warnings.warn(  # pragma: no cover
-                "Failed to identify an unmarshaller for the associated type-variable pair: "
-                f"Original ref: {fkey}, Resolved ref: {resolved}. Will default to no-op.",
-                stacklevel=3,
-            )
-            fields_by_var[name] = NoOpMarshaller(hint, self.context, var=name)
+            fields_by_var[name] = m
+
         return fields_by_var
 
     def __call__(self, val: _ST) -> MarshalledMappingT:

src/typelib/py/refs.py

@@ -18,15 +18,15 @@
 import sys
 import typing
 
-from typelib.py import frames, future
+from typelib.py import frames, future, inspection
 
 __all__ = ("ForwardRef", "evaluate", "forwardref")
 
 ForwardRef: typing.TypeAlias = typing.ForwardRef
 
 
 def forwardref(
-    ref: str,
+    ref: str | type,
     *,
     is_argument: bool = False,
     module: typing.Any | None = None,
@@ -43,11 +43,18 @@ def forwardref(
         module: The python module in which the reference string is defined (optional)
         is_class: Whether the reference string is a class (default True).
     """
+    if not isinstance(ref, str):
+        name = inspection.qualname(ref)
+        module = module or getattr(ref, "__module__", None)
+    else:
+        name = typing.cast(str, ref)
+
     module = _resolve_module_name(ref, module)
     if module is not None:
-        ref = ref.replace(f"{module}.", "")
+        name = name.replace(f"{module}.", "")
+
     return ForwardRef(
-        ref,
+        name,
         is_argument=is_argument,
         module=module,
         is_class=is_class,

src/typelib/unmarshals/api.py

@@ -40,26 +40,24 @@ def unmarshaller(
              May be a type, type alias, [`typing.ForwardRef`][], or string reference.
     """
     nodes = graph.static_order(t)
-    context: dict[type | graph.TypeNode, routines.AbstractUnmarshaller] = (
-        ctx.TypeContext()
-    )
+    context: ctx.TypeContext[routines.AbstractUnmarshaller] = ctx.TypeContext()
     if not nodes:
         return routines.NoOpUnmarshaller(t=t, context=context, var=None)  # type: ignore[arg-type]
 
     # "root" type will always be the final node in the sequence.
     root = nodes[-1]
     for node in nodes:
-        context[node] = _get_unmarshaller(node, context=context)
+        context[node.type] = _get_unmarshaller(node, context=context)
 
-    return context[root]
+    return context[root.type]
 
 
 def _get_unmarshaller(  # type: ignore[return]
     node: graph.TypeNode,
     context: routines.ContextT,
 ) -> routines.AbstractUnmarshaller[T]:
-    if node in context:
-        return context[node]
+    if node.type in context:
+        return context[node.type]
 
     for check, unmarshaller_cls in _HANDLERS.items():
         if check(node.type):

src/typelib/unmarshals/routines.py

@@ -16,7 +16,7 @@
 import uuid
 import warnings
 
-from typelib import constants, graph, serdes
+from typelib import constants, ctx, serdes
 from typelib.py import compat, inspection, refs
 
 T = tp.TypeVar("T")
@@ -134,7 +134,7 @@ def __call__(self, val: tp.Any) -> None:
         return None
 
 
-ContextT: tp.TypeAlias = "tp.Mapping[type | graph.TypeNode, AbstractUnmarshaller]"
+ContextT: tp.TypeAlias = "ctx.TypeContext[AbstractUnmarshaller]"
 BytesT = tp.TypeVar("BytesT", bound=bytes)
 
 
@@ -977,25 +977,21 @@ def __init__(self, t: type[_ST], context: ContextT, *, var: str | None = None):
 
     def _fields_by_var(self):
         fields_by_var = {}
-        tp_var_map = {(t.type, t.var): m for t, m in self.context.items()}
         hints = inspection.cached_type_hints(self.t)
         for name, hint in hints.items():
             resolved = refs.evaluate(hint)
-            fkey = (hint, name)
-            rkey = (resolved, name)
-            if fkey in tp_var_map:
-                fields_by_var[name] = tp_var_map[fkey]
-                continue
-            if rkey in tp_var_map:
-                fields_by_var[name] = tp_var_map[rkey]
+            m = self.context.get(hint) or self.context.get(resolved)
+            if m is None:
+                warnings.warn(
+                    "Failed to identify an unmarshaller for the associated type-variable pair: "
+                    f"Original ref: {hint}, Resolved ref: {resolved}. Will default to no-op.",
+                    stacklevel=4,
+                )
+                fields_by_var[name] = NoOpUnmarshaller(hint, self.context, var=name)
                 continue
 
-            warnings.warn(
-                "Failed to identify an unmarshaller for the associated type-variable pair: "
-                f"Original ref: {fkey}, Resolved ref: {resolved}. Will default to no-op.",
-                stacklevel=3,
-            )
-            fields_by_var[name] = NoOpUnmarshaller(hint, self.context, var=name)
+            fields_by_var[name] = m
+
         return fields_by_var
 
     def __call__(self, val: tp.Any) -> _ST:

tests/unit/unmarshals/test_routines.py

@@ -10,7 +10,6 @@
 
 import pytest
 
-from typelib import graph
 from typelib.unmarshals import routines
 
 from tests import models
@@ -739,12 +738,8 @@ def test_fixed_tuple_unmarshaller(
 @pytest.mark.suite(
     context=dict(
         given_context={
-            graph.TypeNode(int, var="value"): routines.NumberUnmarshaller(
-                int, {}, var="value"
-            ),
-            graph.TypeNode(str, var="field"): routines.StringUnmarshaller(
-                str, {}, var="field"
-            ),
+            int: routines.NumberUnmarshaller(int, {}, var="value"),
+            str: routines.StringUnmarshaller(str, {}, var="field"),
         },
     ),
 )