diff --git a/dimos/core/_test_future_annotations_helper.py b/dimos/core/_test_future_annotations_helper.py new file mode 100644 index 0000000000..a5fe8e6e37 --- /dev/null +++ b/dimos/core/_test_future_annotations_helper.py @@ -0,0 +1,36 @@ +# Copyright 2025 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Helper module for testing blueprint handling with PEP 563 (future annotations). + +This file exists because `from __future__ import annotations` affects the entire file. +""" + +from __future__ import annotations + +from dimos.core.module import Module +from dimos.core.stream import In, Out # noqa + + +class FutureData: + pass + + +class FutureModuleOut(Module): + data: Out[FutureData] = None # type: ignore[assignment] + + +class FutureModuleIn(Module): + data: In[FutureData] = None # type: ignore[assignment] diff --git a/dimos/core/blueprints.py b/dimos/core/blueprints.py index 525c2de42a..484e076af9 100644 --- a/dimos/core/blueprints.py +++ b/dimos/core/blueprints.py @@ -21,7 +21,7 @@ import operator import sys from types import MappingProxyType -from typing import Any, Literal, get_args, get_origin +from typing import Any, Literal, get_args, get_origin, get_type_hints from dimos.core.global_config import GlobalConfig from dimos.core.module import Module @@ -293,16 +293,21 @@ def _make_module_blueprint( ) -> ModuleBlueprint: connections: list[ModuleConnection] = [] - all_annotations = {} - for base_class in reversed(module.__mro__): - if hasattr(base_class, "__annotations__"): - all_annotations.update(base_class.__annotations__) + # Use get_type_hints() to properly resolve string annotations. + try: + all_annotations = get_type_hints(module) + except Exception: + # Fallback to raw annotations if get_type_hints fails. + all_annotations = {} + for base_class in reversed(module.__mro__): + if hasattr(base_class, "__annotations__"): + all_annotations.update(base_class.__annotations__) for name, annotation in all_annotations.items(): origin = get_origin(annotation) - if origin not in (In, Out): # type: ignore[comparison-overlap] + if origin not in (In, Out): continue - direction = "in" if origin == In else "out" # type: ignore[comparison-overlap] + direction = "in" if origin == In else "out" type_ = get_args(annotation)[0] connections.append(ModuleConnection(name=name, type=type_, direction=direction)) # type: ignore[arg-type] diff --git a/dimos/core/test_blueprints.py b/dimos/core/test_blueprints.py index 373a6b8f6e..1ed93dea89 100644 --- a/dimos/core/test_blueprints.py +++ b/dimos/core/test_blueprints.py @@ -14,6 +14,11 @@ import pytest +from dimos.core._test_future_annotations_helper import ( + FutureData, + FutureModuleIn, + FutureModuleOut, +) from dimos.core.blueprints import ( ModuleBlueprint, ModuleBlueprintSet, @@ -317,3 +322,50 @@ class TargetModule(Module): finally: coordinator.stop() + + +def test_future_annotations_support() -> None: + """Test that modules using `from __future__ import annotations` work correctly. + + PEP 563 (future annotations) stores annotations as strings instead of actual types. + This test verifies that _make_module_blueprint properly resolves string annotations + to the actual In/Out types. + """ + + # Test that connections are properly extracted from modules with future annotations + out_blueprint = _make_module_blueprint(FutureModuleOut, args=(), kwargs={}) + assert len(out_blueprint.connections) == 1 + assert out_blueprint.connections[0] == ModuleConnection( + name="data", type=FutureData, direction="out" + ) + + in_blueprint = _make_module_blueprint(FutureModuleIn, args=(), kwargs={}) + assert len(in_blueprint.connections) == 1 + assert in_blueprint.connections[0] == ModuleConnection( + name="data", type=FutureData, direction="in" + ) + + +def test_future_annotations_autoconnect() -> None: + """Test that autoconnect works with modules using `from __future__ import annotations`.""" + + blueprint_set = autoconnect(FutureModuleOut.blueprint(), FutureModuleIn.blueprint()) + + coordinator = blueprint_set.build(GlobalConfig()) + + try: + out_instance = coordinator.get_instance(FutureModuleOut) + in_instance = coordinator.get_instance(FutureModuleIn) + + assert out_instance is not None + assert in_instance is not None + + # Both should have transports set + assert out_instance.data.transport is not None + assert in_instance.data.transport is not None + + # They should be connected via the same transport + assert out_instance.data.transport.topic == in_instance.data.transport.topic + + finally: + coordinator.stop()