Allow signature and property to match in either direction (#151)

Signed-off-by: Tim Paine <3105306+timkpaine@users.noreply.github.com>

Author: timkpaine

ccflow/callable.py

@@ -110,6 +110,7 @@ def _check_signature(self):
                 not isinstance(type_call_arg, TypeVar)
                 and type_call_arg is not Signature.empty
                 and (not isclass(type_call_arg) or not issubclass(type_call_arg, self.context_type))
+                and (not isclass(self.context_type) or not issubclass(self.context_type, type_call_arg))
             ):
                 err_msg_type_mismatch = (
                     f"The context_type {self.context_type} must match the type of the context accepted by __call__ {type_call_arg}"
@@ -122,6 +123,7 @@ def _check_signature(self):
                 not isinstance(type_call_return, TypeVar)
                 and type_call_return is not Signature.empty
                 and (not isclass(type_call_return) or not issubclass(type_call_return, self.result_type))
+                and (not isclass(self.result_type) or not issubclass(self.result_type, type_call_return))
             ):
                 err_msg_type_mismatch = f"The result_type {self.result_type} must match the return type of __call__ {type_call_return}"
                 raise ValueError(err_msg_type_mismatch)

ccflow/tests/test_callable.py

@@ -1,5 +1,5 @@
 from pickle import dumps as pdumps, loads as ploads
-from typing import Generic, List, Optional, TypeVar
+from typing import Generic, List, Optional, Tuple, Type, TypeVar
 from unittest import TestCase
 
 import ray
@@ -268,6 +268,30 @@ def __call__(self, context: NullContext) -> GenericResult[int]:
         return GenericResult[int](value=42)
 
 
+class MyGenericContext(ContextBase, Generic[TContext]):
+    value: TContext
+
+
+class ModelMixedGenericsEnforceContextMatch(CallableModel, Generic[TContext, TResult]):
+    model: CallableModelGenericType[TContext, TResult]
+
+    @property
+    def context_type(self) -> Type[ContextType]:
+        return MyGenericContext[self.model.context_type]
+
+    @property
+    def result_type(self) -> Type[ResultType]:
+        return GenericResult[self.model.result_type]
+
+    @Flow.deps
+    def __deps__(self, context: MyGenericContext[TContext]) -> List[Tuple[CallableModelGenericType[TContext, TResult], List[ContextType]]]:
+        return []
+
+    @Flow.call
+    def __call__(self, context: MyGenericContext[TContext]) -> TResult:
+        return GenericResult(value=None)
+
+
 class TestContext(TestCase):
     def test_immutable(self):
         x = MyContext(a="foo")
@@ -413,6 +437,10 @@ def test_identity(self):
         self.assertEqual(ident(context), context)
         self.assertIsNot(ident(context), context)
 
+    def test_context_call_match_enforcement_generic_base(self):
+        # This should not raise
+        _ = ModelMixedGenericsEnforceContextMatch(model=IdentityCallable())
+
 
 class TestWrapperModel(TestCase):
     def test_wrapper(self):