mypy

Author: stellasia

tests/unit/experimental/pipeline/components.py

@@ -45,6 +45,9 @@ async def run(self, number1: int, number2: int = 2) -> IntResultModel:
 
 
 class ComponentMultiplyWithContext(Component):
+    async def run(self, number1: int, number2: int) -> IntResultModel:
+        return IntResultModel(result=number1 * number2)
+
     async def run_with_context(
         self, context_: RunContext, number1: int, number2: int = 2
     ) -> IntResultModel:

tests/unit/experimental/pipeline/test_component.py

@@ -41,7 +41,9 @@ async def test_component_run() -> None:
     result = await c.run(number1=1, number2=2)
     assert isinstance(result, IntResultModel)
     assert isinstance(
-        result.result, ComponentMultiply.component_outputs["result"]["annotation"]
+        result.result,
+        # we know this is a type and not a bool or str:
+        ComponentMultiply.component_outputs["result"]["annotation"],  # type: ignore
     )
 
 
@@ -55,7 +57,9 @@ async def test_component_run_with_context_default_implementation() -> None:
         number1=1,
         number2=2,
     )
-    assert result.result == 2
+    # the type checker doesn't know about the type
+    # because the method is not re-declared
+    assert result.result == 2  # type: ignore
 
 
 @pytest.mark.asyncio