PR comments

Author: tconley1428

scripts/gen_visitors.py scripts/gen_payload_visitor.pyscripts/gen_visitors.py renamed to scripts/gen_payload_visitor.py

@@ -39,18 +39,18 @@ def emit_loop(
 
 
 def emit_singular(
-    field_name: str, access_expr: str, child_method: str, check_presence: bool
+    field_name: str, access_expr: str, child_method: str, presence_word: Optional[str]
 ) -> str:
     # Helper to emit a singular field visit with presence check and optional headers guard
-    if check_presence:
+    if presence_word:
         if field_name == "headers":
             return f"""\
         if not self.skip_headers:
-            if o.HasField("{field_name}"):
+            {presence_word} o.HasField("{field_name}"):
                 await self._visit_{child_method}(fs, {access_expr})"""
         else:
             return f"""\
-        if o.HasField("{field_name}"):
+        {presence_word} o.HasField("{field_name}"):
             await self._visit_{child_method}(fs, {access_expr})"""
     else:
         if field_name == "headers":
@@ -63,6 +63,67 @@ def emit_singular(
 
 
 class VisitorGenerator:
+    def generate(self, roots: list[Descriptor]) -> str:
+        """
+        Generate Python source code that, given a function f(Payload) -> Payload,
+        applies it to every Payload contained within a WorkflowActivation tree.
+
+        The generated code defines async visitor functions for each reachable
+        protobuf message type starting from WorkflowActivation, including support
+        for repeated fields and map entries, and a convenience entrypoint
+        function `visit`.
+        """
+
+        for r in roots:
+            self.walk(r)
+
+        header = """
+# This file is generated by gen_payload_visitor.py. Changes should be made there.
+import abc
+from typing import Any, MutableSequence
+
+from temporalio.api.common.v1.message_pb2 import Payload
+
+class VisitorFunctions(abc.ABC):
+    \"\"\"Set of functions which can be called by the visitor. 
+    Allows handling payloads as a sequence.
+    \"\"\"
+    @abc.abstractmethod
+    async def visit_payload(self, payload: Payload) -> None:
+        \"\"\"Called when encountering a single payload.\"\"\"
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None:
+        \"\"\"Called when encountering multiple payloads together.\"\"\"
+        raise NotImplementedError()
+
+class PayloadVisitor:
+    \"\"\"A visitor for payloads. 
+    Applies a function to every payload in a tree of messages.
+    \"\"\"
+    def __init__(
+        self, *, skip_search_attributes: bool = False, skip_headers: bool = False
+    ):
+        \"\"\"Creates a new payload visitor.\"\"\"
+        self.skip_search_attributes = skip_search_attributes
+        self.skip_headers = skip_headers
+
+    async def visit(
+        self, fs: VisitorFunctions, root: Any
+    ) -> None:
+        \"\"\"Visits the given root message with the given function.\"\"\"
+        method_name = "_visit_" + root.DESCRIPTOR.full_name.replace(".", "_")
+        method = getattr(self, method_name, None)
+        if method is not None:
+            await method(fs, root)
+        else:
+            raise ValueError(f"Unknown root message type: {root.DESCRIPTOR.full_name}")
+
+"""
+
+        return header + "\n".join(self.methods)
+
     def __init__(self):
         # Track which message descriptors have visitor methods generated
         self.generated: dict[str, bool] = {
@@ -71,21 +132,24 @@ def __init__(self):
         }
         self.in_progress: set[str] = set()
         self.methods: list[str] = [
-            """    async def _visit_temporal_api_common_v1_Payload(self, fs, o):
-            await fs.visit_payload(o)
+            """\
+    async def _visit_temporal_api_common_v1_Payload(self, fs, o):
+        await fs.visit_payload(o)
     """,
-            """    async def _visit_temporal_api_common_v1_Payloads(self, fs, o):
-            await fs.visit_payloads(o.payloads)
+            """\
+    async def _visit_temporal_api_common_v1_Payloads(self, fs, o):
+        await fs.visit_payloads(o.payloads)
     """,
-            """    async def _visit_payload_container(self, fs, o):
-            await fs.visit_payloads(o)
+            """\
+    async def _visit_payload_container(self, fs, o):
+        await fs.visit_payloads(o)
     """,
         ]
 
     def check_repeated(self, child_desc, field, iter_expr) -> Optional[str]:
         # Special case for repeated payloads, handle them directly
         if child_desc.full_name == Payload.DESCRIPTOR.full_name:
-            return emit_singular(field.name, iter_expr, "payload_container", False)
+            return emit_singular(field.name, iter_expr, "payload_container", None)
         else:
             child_needed = self.walk(child_desc)
             if child_needed:
@@ -105,7 +169,7 @@ def walk(self, desc: Descriptor) -> bool:
             # Break cycles; if another path proves this node needed, we'll revisit
             return False
 
-        needed = False
+        has_payload = False
         self.in_progress.add(key)
         lines: list[str] = [f"    async def _visit_{name_for(desc)}(self, fs, o):"]
         # If this is the SearchAttributes message, allow skipping
@@ -146,7 +210,7 @@ def walk(self, desc: Descriptor) -> bool:
                         child_desc = val_fd.message_type
                         child_needed = self.walk(child_desc)
                         if child_needed:
-                            needed = True
+                            has_payload = True
                             lines.append(
                                 emit_loop(
                                     field.name,
@@ -163,7 +227,7 @@ def walk(self, desc: Descriptor) -> bool:
                         child_desc = key_fd.message_type
                         child_needed = self.walk(child_desc)
                         if child_needed:
-                            needed = True
+                            has_payload = True
                             lines.append(
                                 emit_loop(
                                     field.name,
@@ -176,16 +240,16 @@ def walk(self, desc: Descriptor) -> bool:
                         field.message_type, field, f"o.{field.name}"
                     )
                     if child is not None:
-                        needed = True
+                        has_payload = True
                         lines.append(child)
             else:
                 child_desc = field.message_type
-                child_needed = self.walk(child_desc)
-                needed |= child_needed
-                if child_needed:
+                child_has_payload = self.walk(child_desc)
+                has_payload |= child_has_payload
+                if child_has_payload:
                     lines.append(
                         emit_singular(
-                            field.name, f"o.{field.name}", name_for(child_desc), True
+                            field.name, f"o.{field.name}", name_for(child_desc), "if"
                         )
                     )
 
@@ -195,81 +259,23 @@ def walk(self, desc: Descriptor) -> bool:
             first = True
             for field in fields:
                 child_desc = field.message_type
-                child_needed = self.walk(child_desc)
-                needed |= child_needed
-                if child_needed:
+                child_has_payload = self.walk(child_desc)
+                has_payload |= child_has_payload
+                if child_has_payload:
                     if_word = "if" if first else "elif"
                     first = False
                     line = emit_singular(
-                        field.name, f"o.{field.name}", name_for(child_desc), True
-                    ).replace("        if", f"        {if_word}", 1)
+                        field.name, f"o.{field.name}", name_for(child_desc), if_word
+                    )
                     oneof_lines.append(line)
             if oneof_lines:
                 lines.extend(oneof_lines)
 
-        self.generated[key] = needed
+        self.generated[key] = has_payload
         self.in_progress.discard(key)
-        if needed:
+        if has_payload:
             self.methods.append("\n".join(lines) + "\n")
-        return needed
-
-    def generate(self, roots: list[Descriptor]) -> str:
-        """
-        Generate Python source code that, given a function f(Payload) -> Payload,
-        applies it to every Payload contained within a WorkflowActivation tree.
-
-        The generated code defines async visitor functions for each reachable
-        protobuf message type starting from WorkflowActivation, including support
-        for repeated fields and map entries, and a convenience entrypoint
-        function `visit`.
-        """
-
-        # We avoid importing google.api deps in service protos; expand by walking from
-        # WorkflowActivationCompletion root which references many command messages.
-        for r in roots:
-            self.walk(r)
-
-        header = """
-import abc
-from typing import Any, MutableSequence
-
-from temporalio.api.common.v1.message_pb2 import Payload
-
-class VisitorFunctions(abc.ABC):
-    \"\"\"Set of functions which can be called by the visitor. Allows handling payloads as a sequence.\"\"\"
-    @abc.abstractmethod
-    async def visit_payload(self, payload: Payload) -> None:
-        \"\"\"Called when encountering a single payload.\"\"\"
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None:
-        \"\"\"Called when encountering multiple payloads together.\"\"\"
-        raise NotImplementedError()
-
-class PayloadVisitor:
-    \"\"\"A visitor for payloads. Applies a function to every payload in a tree of messages.\"\"\"
-    def __init__(
-        self, *, skip_search_attributes: bool = False, skip_headers: bool = False
-    ):
-        \"\"\"Creates a new payload visitor.\"\"\"
-        self.skip_search_attributes = skip_search_attributes
-        self.skip_headers = skip_headers
-
-    async def visit(
-        self, fs: VisitorFunctions, root: Any
-    ) -> None:
-        \"\"\"Visits the given root message with the given function.\"\"\"
-        method_name = "_visit_" + root.DESCRIPTOR.full_name.replace(".", "_")
-        method = getattr(self, method_name, None)
-        if method is not None:
-            await method(fs, root)
-        else:
-            raise ValueError(f"Unknown root message type: {root.DESCRIPTOR.full_name}")
-
-"""
-
-        return header + "\n".join(self.methods)
+        return has_payload
 
 
 def write_generated_visitors_into_visitor_generated_py() -> None:

scripts/gen_protos_docker.py

@@ -33,5 +33,6 @@
 subprocess.run(["uv", "run", "poe", "format"], check=True)
 
 subprocess.run(
-    ["uv", "run", os.path.join(os.getcwd(), "scripts", "gen_visitors.py")], check=True
+    ["uv", "run", os.path.join(os.getcwd(), "scripts", "gen_payload_visitor.py")],
+    check=True,
 )

temporalio/bridge/_visitor.py

@@ -1,11 +1,14 @@
+# This file is generated by gen_payload_visitor.py. Changes should be made there.
 import abc
 from typing import Any, MutableSequence
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
 
 class VisitorFunctions(abc.ABC):
-    """Set of functions which can be called by the visitor. Allows handling payloads as a sequence."""
+    """Set of functions which can be called by the visitor.
+    Allows handling payloads as a sequence.
+    """
 
     @abc.abstractmethod
     async def visit_payload(self, payload: Payload) -> None:
@@ -19,7 +22,9 @@ async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None:
 
 
 class PayloadVisitor:
-    """A visitor for payloads. Applies a function to every payload in a tree of messages."""
+    """A visitor for payloads.
+    Applies a function to every payload in a tree of messages.
+    """
 
     def __init__(
         self, *, skip_search_attributes: bool = False, skip_headers: bool = False