Merge pull request #112 from timkpaine/tkp/pub

Release publisher model, add tests

Author: timkpaine

ccflow/__init__.py

@@ -6,6 +6,7 @@
 from .context import *
 from .enums import Enum
 from .exttypes import *
+from .models import *
 from .object_config import *
 from .publisher import *
 from .result import *

ccflow/models/__init__.py

@@ -1 +1 @@
-
+from .publisher import *

ccflow/models/publisher.py

@@ -0,0 +1,60 @@
+from typing import Generic, Type
+
+from pydantic import Field
+from typing_extensions import override
+
+from ..callable import CallableModelType, ContextType, Flow, ResultType, WrapperModel
+from ..publisher import PublisherType
+from ..result import GenericResult
+
+__all__ = ("PublisherModel",)
+
+
+class PublisherModel(
+    WrapperModel[CallableModelType],
+    Generic[CallableModelType, PublisherType],
+):
+    """Model that chains together a callable model and a publisher to publish the results of the callable model."""
+
+    publisher: PublisherType
+    field: str = Field(None, description="Specific field on model output to publish")
+    return_data: bool = Field(
+        False,
+        description="Whether to return the underlying model result as the output instead of the publisher output",
+    )
+
+    @property
+    def result_type(self) -> Type[ResultType]:
+        """Result type that will be returned. Could be over-ridden by child class."""
+        if self.return_data:
+            return self.model.result_type
+        else:
+            return GenericResult
+
+    def _get_publisher(self, context):
+        publisher = self.publisher.model_copy()
+        # Set the name, if needed
+        if not publisher.name and self.meta.name:
+            publisher.name = self.meta.name
+        # Augment any existing name parameters with the context parameters
+        name_params = publisher.name_params.copy()
+        name_params.update(context.model_dump(exclude={"type_"}))
+        publisher.name_params = name_params
+        return publisher
+
+    @override
+    @Flow.call
+    def __call__(self, context: ContextType) -> ResultType:
+        """This method gets the result from the underlying model, and publishes it."""
+        publisher = self._get_publisher(context)
+        data = self.model(context)
+        if self.field:
+            pub_data = getattr(data, self.field)
+        else:
+            pub_data = data
+        publisher.data = pub_data
+        out = publisher()
+        if self.return_data:
+            return data
+        else:
+            return self.result_type(value=out)

ccflow/tests/models/__init__.py

@@ -0,0 +1,23 @@
+from unittest.mock import patch
+
+from ccflow import CallableModel, DictResult, Flow, GenericResult, NullContext
+from ccflow.models import PublisherModel
+from ccflow.publishers import PrintPublisher
+
+
+class TestModel(CallableModel):
+    @Flow.call
+    def __call__(self, context: NullContext) -> DictResult[str, str]:
+        return DictResult[str, str](value={"message": "Hello, World!"})
+
+
+class TestPublisherModel:
+    def test_run(self):
+        with patch("ccflow.publishers.print.print") as mock_print:
+            model = PublisherModel(model=TestModel(), publisher=PrintPublisher())
+            res = model(None)
+            assert isinstance(res, GenericResult)  # from PrintPublisher
+            assert isinstance(res.value, DictResult[str, str])
+            assert res.value.value == {"message": "Hello, World!"}
+            assert mock_print.call_count == 1
+            assert mock_print.call_args[0][0].value == {"message": "Hello, World!"}