Resolve generic stream types in graph metadata

src/ezmsg/core/backend.py

@@ -14,6 +14,7 @@
 from multiprocessing.synchronize import Barrier as BarrierType
 from multiprocessing.connection import wait, Connection
 from socket import socket
+from typing import get_origin
 
 from .netprotocol import DEFAULT_SHM_SIZE, AddressType
 
@@ -31,6 +32,7 @@
     OutputRelay,
 )
 from .unit import Unit, PROCESS_ATTR, SUBSCRIBES_ATTR, PUBLISHES_ATTR
+from .type_resolution import resolve_stream_type
 from .settings import Settings
 from .graphmeta import (
     CollectionMetadata,
@@ -391,7 +393,9 @@ def _type_name(self, tp: type) -> str:
         return f"{tp.__module__}.{tp.__qualname__}"
 
     def _stream_type_name(self, stream_type: object) -> str:
-        if inspect.isclass(stream_type):
+        # Python 3.10 considers GenericAlias objects (e.g. list[int]) classes.
+        # Preserve their parameters instead of naming only the origin class.
+        if get_origin(stream_type) is None and inspect.isclass(stream_type):
             return self._type_name(stream_type)
         return repr(stream_type)
 
@@ -421,7 +425,9 @@ def _component_metadata(self) -> GraphMetadata:
                         else None
                     ),
                     settings_type=(
-                        self._stream_type_name(input_settings.msg_type)
+                        self._stream_type_name(
+                            resolve_stream_type(type(comp), input_settings.msg_type)
+                        )
                         if isinstance(input_settings, InputStream)
                         else None
                     ),
@@ -431,7 +437,9 @@ def _component_metadata(self) -> GraphMetadata:
                 topic_entries: dict[str, TopicMetadataType] = {}
                 relay_entries: dict[str, RelayMetadataType] = {}
                 for stream_name, stream in comp.streams.items():
-                    msg_type = self._stream_type_name(stream.msg_type)
+                    msg_type = self._stream_type_name(
+                        resolve_stream_type(type(comp), stream.msg_type)
+                    )
                     if isinstance(stream, InputRelay):
                         runtime = _relay_runtime_info(stream)
                         relay_entries[stream_name] = InputRelayMetadata(

src/ezmsg/core/type_resolution.py

@@ -0,0 +1,55 @@
+"""Resolve inherited generic annotations without changing shared stream objects."""
+
+import types
+from functools import reduce
+from operator import or_
+from typing import TypeVar, Union, get_args, get_origin
+
+
+def _substitute(annotation, bindings):
+    if isinstance(annotation, TypeVar):
+        seen = set()
+        while isinstance(annotation, TypeVar) and annotation in bindings and annotation not in seen:
+            seen.add(annotation)
+            annotation = bindings[annotation]
+        return annotation
+    args = get_args(annotation)
+    if not args:
+        return annotation
+    resolved = tuple(_substitute(arg, bindings) for arg in args)
+    if resolved == args:
+        return annotation
+    if hasattr(annotation, "copy_with"):
+        return annotation.copy_with(resolved)
+    origin = get_origin(annotation)
+    if origin in (Union, types.UnionType):
+        return reduce(or_, resolved)
+    try:
+        return origin[resolved[0] if len(resolved) == 1 else resolved]
+    except (TypeError, AttributeError):
+        return annotation
+
+
+def resolve_stream_type(component_type: type, annotation):
+    """Substitute explicit generic base arguments; leave unbound types unresolved."""
+    candidates = {}
+
+    def walk(cls, inherited):
+        for parent in cls.__dict__.get("__orig_bases__", cls.__bases__):
+            origin = get_origin(parent) or parent
+            args = tuple(_substitute(arg, inherited) for arg in get_args(parent))
+            bindings = {**inherited, **dict(zip(getattr(origin, "__parameters__", ()), args))}
+            for variable, value in bindings.items():
+                if variable != value:
+                    candidates.setdefault(variable, []).append(value)
+            if isinstance(origin, type) and origin is not object:
+                walk(origin, bindings)
+
+    walk(component_type, {})
+    # Ambiguous multiple inheritance must not advertise an invented concrete type.
+    bindings = {
+        variable: values[0]
+        for variable, values in candidates.items()
+        if all(value == values[0] for value in values)
+    }
+    return _substitute(annotation, bindings)

tests/test_stream_type_resolution.py

@@ -0,0 +1,68 @@
+from typing import Generic, TypeVar
+
+import pytest
+
+import ezmsg.core as ez
+from ezmsg.core.backend import ExecutionContext, GraphRunner
+from ezmsg.core.type_resolution import resolve_stream_type
+
+T = TypeVar("T")
+U = TypeVar("U")
+S = TypeVar("S")
+
+
+class GenericUnit(ez.Unit, Generic[S, T, U]):
+    INPUT_SETTINGS = ez.InputStream(S)
+    INPUT = ez.InputStream(T)
+    OUTPUT = ez.OutputStream(U)
+
+
+class Intermediate(GenericUnit[ez.Settings, T, list[T]], Generic[T]):
+    pass
+
+
+class Concrete(Intermediate[int]):
+    pass
+
+
+def test_specialized_stream_and_settings_metadata():
+    unit = Concrete()
+    ExecutionContext.setup({"UNIT": unit})
+    metadata = GraphRunner(components={"UNIT": unit})._component_metadata().components["UNIT"]
+    assert metadata.streams["INPUT"].msg_type == "builtins.int"
+    assert metadata.streams["OUTPUT"].msg_type == "list[int]"
+    assert metadata.dynamic_settings.settings_type == "ezmsg.core.settings.Settings"
+    assert GenericUnit.__streams__["INPUT"].msg_type is T
+    assert unit.INPUT.msg_type is T  # Metadata does not mutate stream declarations.
+
+
+def test_unspecialized_variables_remain_unresolved():
+    assert resolve_stream_type(GenericUnit, T) is T
+
+
+def test_independent_specializations():
+    class Text(Intermediate[str]):
+        pass
+
+    assert resolve_stream_type(Text, U) == list[str]
+    assert resolve_stream_type(Concrete, U) == list[int]
+
+
+@pytest.mark.parametrize(
+    "annotation, expected",
+    [
+        (int, "builtins.int"),
+        (list, "builtins.list"),
+        (list[int], "list[int]"),
+        (dict[str, list[int]], "dict[str, list[int]]"),
+        (tuple[int, str], "tuple[int, str]"),
+    ],
+)
+def test_parameterized_stream_metadata_preserves_arguments(annotation, expected):
+    class TypedUnit(ez.Unit):
+        OUTPUT = ez.OutputStream(annotation)
+
+    unit = TypedUnit()
+    ExecutionContext.setup({"UNIT": unit})
+    metadata = GraphRunner(components={"UNIT": unit})._component_metadata().components["UNIT"]
+    assert metadata.streams["OUTPUT"].msg_type == expected