Copy of pydantic v1 converter from samples-python

Author: dandavison

temporalio/contrib/pydantic/converter.py

@@ -0,0 +1,58 @@
+import json
+from typing import Any, Optional
+
+from pydantic.json import pydantic_encoder
+from temporalio.api.common.v1 import Payload
+from temporalio.converter import (
+    CompositePayloadConverter,
+    DataConverter,
+    DefaultPayloadConverter,
+    JSONPlainPayloadConverter,
+)
+
+
+class PydanticJSONPayloadConverter(JSONPlainPayloadConverter):
+    """Pydantic JSON payload converter.
+
+    This extends the :py:class:`JSONPlainPayloadConverter` to override
+    :py:meth:`to_payload` using the Pydantic encoder.
+    """
+
+    def to_payload(self, value: Any) -> Optional[Payload]:
+        """Convert all values with Pydantic encoder or fail.
+
+        Like the base class, we fail if we cannot convert. This payload
+        converter is expected to be the last in the chain, so it can fail if
+        unable to convert.
+        """
+        # We let JSON conversion errors be thrown to caller
+        return Payload(
+            metadata={"encoding": self.encoding.encode()},
+            data=json.dumps(
+                value, separators=(",", ":"), sort_keys=True, default=pydantic_encoder
+            ).encode(),
+        )
+
+
+class PydanticPayloadConverter(CompositePayloadConverter):
+    """Payload converter that replaces Temporal JSON conversion with Pydantic
+    JSON conversion.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(
+            *(
+                (
+                    c
+                    if not isinstance(c, JSONPlainPayloadConverter)
+                    else PydanticJSONPayloadConverter()
+                )
+                for c in DefaultPayloadConverter.default_encoding_payload_converters
+            )
+        )
+
+
+pydantic_data_converter = DataConverter(
+    payload_converter_class=PydanticPayloadConverter
+)
+"""Data converter using Pydantic JSON conversion."""

tests/contrib/test_pydantic.py

@@ -0,0 +1,94 @@
+import dataclasses
+import uuid
+from datetime import datetime, timedelta
+from ipaddress import IPv4Address
+from typing import List
+
+from pydantic import BaseModel
+
+from temporalio import activity, workflow
+from temporalio.client import Client
+from temporalio.contrib.pydantic.converter import pydantic_data_converter
+from temporalio.worker import Worker
+from temporalio.worker.workflow_sandbox import (
+    SandboxedWorkflowRunner,
+    SandboxRestrictions,
+)
+
+
+class MyPydanticModel(BaseModel):
+    some_ip: IPv4Address
+    some_date: datetime
+
+
+@activity.defn
+async def my_activity(models: List[MyPydanticModel]) -> List[MyPydanticModel]:
+    activity.logger.info("Got models in activity: %s" % models)
+    return models
+
+
+@workflow.defn
+class MyWorkflow:
+    @workflow.run
+    async def run(self, models: List[MyPydanticModel]) -> List[MyPydanticModel]:
+        workflow.logger.info("Got models in workflow: %s" % models)
+        return await workflow.execute_activity(
+            my_activity, models, start_to_close_timeout=timedelta(minutes=1)
+        )
+
+
+# Due to known issues with Pydantic's use of issubclass and our inability to
+# override the check in sandbox, Pydantic will think datetime is actually date
+# in the sandbox. At the expense of protecting against datetime.now() use in
+# workflows, we're going to remove datetime module restrictions. See sdk-python
+# README's discussion of known sandbox issues for more details.
+def new_sandbox_runner() -> SandboxedWorkflowRunner:
+    # TODO(cretz): Use with_child_unrestricted when https://github.com/temporalio/sdk-python/issues/254
+    # is fixed and released
+    invalid_module_member_children = dict(
+        SandboxRestrictions.invalid_module_members_default.children
+    )
+    del invalid_module_member_children["datetime"]
+    return SandboxedWorkflowRunner(
+        restrictions=dataclasses.replace(
+            SandboxRestrictions.default,
+            invalid_module_members=dataclasses.replace(
+                SandboxRestrictions.invalid_module_members_default,
+                children=invalid_module_member_children,
+            ),
+        )
+    )
+
+
+async def test_workflow_with_pydantic_model(client: Client):
+    # Replace data converter in client
+    new_config = client.config()
+    new_config["data_converter"] = pydantic_data_converter
+    client = Client(**new_config)
+    task_queue_name = str(uuid.uuid4())
+
+    orig_models = [
+        MyPydanticModel(
+            some_ip=IPv4Address("127.0.0.1"),
+            some_date=datetime(2000, 1, 2, 3, 4, 5),
+        ),
+        MyPydanticModel(
+            some_ip=IPv4Address("127.0.0.2"),
+            some_date=datetime(2001, 2, 3, 4, 5, 6),
+        ),
+    ]
+
+    async with Worker(
+        client,
+        task_queue=task_queue_name,
+        workflows=[MyWorkflow],
+        activities=[my_activity],
+        workflow_runner=new_sandbox_runner(),
+    ):
+        result = await client.execute_workflow(
+            MyWorkflow.run,
+            orig_models,
+            id=str(uuid.uuid4()),
+            task_queue=task_queue_name,
+        )
+    assert orig_models == result