⚡️ Speed up method WorkflowHandle.get_update_handle_for by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for WorkflowHandle.get_update_handle_for in temporalio/client.py

⏱️ Runtime : 789 microseconds → 681 microseconds (best of 362 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup by eliminating method call overhead and reducing attribute lookups through two key optimizations:

1. Eliminated method call indirection in get_update_handle_for:

• The original version called self.get_update_handle(), adding method call overhead
• The optimized version directly constructs WorkflowUpdateHandle, removing this intermediate call layer
• This eliminates the function call dispatch overhead visible in the profiler (95.7% of time was spent in the method call)

2. Pre-computed expressions to avoid repeated evaluations:

• Moved workflow_run_id or self._run_id logic into a local variable run_id_to_pass
• Cached update._defn.ret_type in local variable ret_type
• This avoids Python's overhead of evaluating the or expression and attribute access during function argument construction

Performance characteristics:

• Most effective for frequently called methods (test cases show 16-28% improvements for basic operations)
• Provides consistent speedup across different input types and sizes
• Particularly beneficial in high-throughput scenarios where these methods are called repeatedly (1000+ handle creations show 16% improvement)
• Error cases also benefit slightly (6-7% faster) due to reduced overhead before exceptions are raised

The optimizations preserve exact functional behavior while reducing Python bytecode overhead, making them especially valuable for performance-critical workflow handle operations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3204 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from typing import Any, Optional, Type

# imports
import pytest  # used for our unit tests
from temporalio.client import WorkflowHandle


# Mocks for temporalio.workflow.UpdateMethodMultiParam and WorkflowUpdateHandle
class DummyUpdateDefn:
    def __init__(self, ret_type: Type):
        self.ret_type = ret_type

class DummyUpdateMethodMultiParam:
    def __init__(self, ret_type: Type):
        self._defn = DummyUpdateDefn(ret_type)

class WorkflowUpdateHandle:
    def __init__(self, client, update_id, workflow_id, workflow_run_id=None, result_type=None):
        self.client = client
        self.update_id = update_id
        self.workflow_id = workflow_id
        self.workflow_run_id = workflow_run_id
        self.result_type = result_type

class DummyClient:
    pass

# --- Unit Tests ---

# Basic Test Cases
def test_basic_handle_creation_with_str_type():
    """Test handle creation with a basic string return type."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow1", run_id="run123")
    update = DummyUpdateMethodMultiParam(str)
    codeflash_output = handle.get_update_handle_for(update, "updateA"); update_handle = codeflash_output # 625ns -> 500ns (25.0% faster)

def test_basic_handle_creation_with_int_type():
    """Test handle creation with an integer return type."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow2", run_id="run456")
    update = DummyUpdateMethodMultiParam(int)
    codeflash_output = handle.get_update_handle_for(update, "updateB"); update_handle = codeflash_output # 541ns -> 459ns (17.9% faster)

def test_handle_creation_with_no_run_id():
    """Test handle creation when no run_id is provided."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow3")
    update = DummyUpdateMethodMultiParam(float)
    codeflash_output = handle.get_update_handle_for(update, "updateC"); update_handle = codeflash_output # 583ns -> 458ns (27.3% faster)

def test_handle_creation_with_explicit_workflow_run_id():
    """Test handle creation with explicit workflow_run_id argument."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow4", run_id="run789")
    update = DummyUpdateMethodMultiParam(dict)
    codeflash_output = handle.get_update_handle_for(update, "updateD", workflow_run_id="run999"); update_handle = codeflash_output # 708ns -> 583ns (21.4% faster)

# Edge Test Cases
def test_handle_creation_with_empty_update_id():
    """Test handle creation with an empty string as update ID."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow5", run_id="run321")
    update = DummyUpdateMethodMultiParam(list)
    codeflash_output = handle.get_update_handle_for(update, ""); update_handle = codeflash_output # 541ns -> 458ns (18.1% faster)


def test_handle_creation_with_none_update_method():
    """Test handle creation with None as update method (should raise AttributeError)."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow7", run_id="run987")
    with pytest.raises(AttributeError):
        handle.get_update_handle_for(None, "updateE") # 667ns -> 625ns (6.72% faster)

def test_handle_creation_with_custom_type():
    """Test handle creation with a custom type as return type."""
    class CustomType:
        pass
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow8", run_id="run111")
    update = DummyUpdateMethodMultiParam(CustomType)
    codeflash_output = handle.get_update_handle_for(update, "updateF"); update_handle = codeflash_output # 750ns -> 583ns (28.6% faster)

def test_handle_creation_with_update_method_missing_defn():
    """Test handle creation with update method missing _defn (should raise AttributeError)."""
    class BadUpdateMethod:
        pass
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow9", run_id="run222")
    bad_update = BadUpdateMethod()
    with pytest.raises(AttributeError):
        handle.get_update_handle_for(bad_update, "updateG") # 625ns -> 625ns (0.000% faster)

def test_handle_creation_with_update_method_defn_missing_ret_type():
    """Test handle creation with update method whose _defn is missing ret_type (should raise AttributeError)."""
    class BadDefn:
        pass
    class BadUpdateMethodMultiParam:
        def __init__(self):
            self._defn = BadDefn()
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow10", run_id="run333")
    bad_update = BadUpdateMethodMultiParam()
    with pytest.raises(AttributeError):
        handle.get_update_handle_for(bad_update, "updateH") # 666ns -> 625ns (6.56% faster)

# Large Scale Test Cases
def test_many_handles_created_unique_ids_and_types():
    """Test creating many handles with unique IDs and types."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow_large", run_id="run_large")
    # Use 1000 handles for scalability test
    for i in range(1000):
        update_id = f"update_{i}"
        # Use alternating types
        ret_type = str if i % 2 == 0 else int
        update = DummyUpdateMethodMultiParam(ret_type)
        codeflash_output = handle.get_update_handle_for(update, update_id); update_handle = codeflash_output # 244μs -> 210μs (16.3% faster)

def test_handles_created_with_large_custom_types():
    """Test creating handles with large custom types."""
    class LargeCustomType:
        def __init__(self, n):
            self.data = [i for i in range(n)]
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow_large_custom", run_id="run_custom")
    # Use 100 custom types
    for i in range(100):
        custom_type = type(f"CustomType{i}", (), {})
        update = DummyUpdateMethodMultiParam(custom_type)
        codeflash_output = handle.get_update_handle_for(update, f"update_custom_{i}"); update_handle = codeflash_output # 25.3μs -> 22.7μs (11.6% faster)

def test_handles_created_with_long_update_ids():
    """Test creating handles with very long update IDs."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow_long_id", run_id="run_long")
    long_id = "x" * 500  # 500-character update ID
    update = DummyUpdateMethodMultiParam(str)
    codeflash_output = handle.get_update_handle_for(update, long_id); update_handle = codeflash_output # 541ns -> 458ns (18.1% faster)

def test_handles_created_with_varied_run_ids():
    """Test creating handles with varied run_ids."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow_varied")
    for i in range(50):
        run_id = f"run_{i}"
        update = DummyUpdateMethodMultiParam(float)
        codeflash_output = handle.get_update_handle_for(update, f"update_varied_{i}", workflow_run_id=run_id); update_handle = codeflash_output # 13.0μs -> 12.4μs (5.03% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

# Patch WorkflowUpdateHandle in WorkflowHandle to use DummyWorkflowUpdateHandle for testing
import sys
from typing import Any, Generic, Optional, Type, Union

# imports
import pytest  # used for our unit tests
from temporalio.client import WorkflowHandle

# --- Mock classes and setup for testing ---

class DummyClient:
    """A minimal stub for the Client, as required by WorkflowHandle."""
    pass

class DummyWorkflowUpdateHandle:
    """A minimal stub for WorkflowUpdateHandle, to check correct construction."""
    def __init__(self, client, update_id, workflow_id, workflow_run_id=None, result_type=None):
        self.client = client
        self.update_id = update_id
        self.workflow_id = workflow_id
        self.workflow_run_id = workflow_run_id
        self.result_type = result_type

# --- Minimal mocks for temporalio.workflow.UpdateMethodMultiParam ---

class DummyUpdateDefn:
    """Mock for ._defn.ret_type property."""
    def __init__(self, ret_type):
        self.ret_type = ret_type

class DummyUpdateMethodMultiParam:
    """Mock for temporalio.workflow.UpdateMethodMultiParam."""
    def __init__(self, ret_type):
        self._defn = DummyUpdateDefn(ret_type)

# --- Unit tests ---

# -------- BASIC TEST CASES --------

def test_basic_handle_returns_correct_type_and_fields():
    """Test that get_update_handle_for returns a DummyWorkflowUpdateHandle with correct fields."""
    client = DummyClient()
    workflow_id = "workflow123"
    run_id = "run456"
    handle = WorkflowHandle(client, workflow_id, run_id=run_id)
    update_method = DummyUpdateMethodMultiParam(ret_type=int)
    update_id = "update789"
    # Call the function
    codeflash_output = handle.get_update_handle_for(update_method, update_id); update_handle = codeflash_output # 542ns -> 542ns (0.000% faster)

def test_basic_handle_with_explicit_workflow_run_id():
    """Test that explicit workflow_run_id overrides handle's run_id."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflowA", run_id="runA")
    update_method = DummyUpdateMethodMultiParam(ret_type=str)
    update_id = "updateB"
    explicit_run_id = "runB"
    codeflash_output = handle.get_update_handle_for(update_method, update_id, workflow_run_id=explicit_run_id); update_handle = codeflash_output # 584ns -> 541ns (7.95% faster)

def test_basic_handle_with_different_return_types():
    """Test that result_type is correctly passed for different types."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    for ret_type in [int, str, float, dict, list]:
        update_method = DummyUpdateMethodMultiParam(ret_type=ret_type)
        codeflash_output = handle.get_update_handle_for(update_method, "update"); update_handle = codeflash_output # 1.62μs -> 1.50μs (8.27% faster)

# -------- EDGE TEST CASES --------

def test_handle_with_none_run_id_and_none_explicit_run_id():
    """Test that workflow_run_id is None if both handle's run_id and explicit workflow_run_id are None."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflowX", run_id=None)
    update_method = DummyUpdateMethodMultiParam(ret_type=bool)
    codeflash_output = handle.get_update_handle_for(update_method, "updateX"); update_handle = codeflash_output # 542ns -> 458ns (18.3% faster)

def test_handle_with_empty_string_ids():
    """Test that empty string IDs are handled correctly."""
    client = DummyClient()
    handle = WorkflowHandle(client, "", run_id="")
    update_method = DummyUpdateMethodMultiParam(ret_type=type(None))
    codeflash_output = handle.get_update_handle_for(update_method, "", workflow_run_id=""); update_handle = codeflash_output # 625ns -> 542ns (15.3% faster)

def test_handle_with_non_string_ids():
    """Test that non-string IDs are accepted and handled (not type checked)."""
    client = DummyClient()
    handle = WorkflowHandle(client, 12345, run_id=67890)
    update_method = DummyUpdateMethodMultiParam(ret_type=list)
    codeflash_output = handle.get_update_handle_for(update_method, 11111); update_handle = codeflash_output # 541ns -> 459ns (17.9% faster)

def test_handle_with_custom_object_return_type():
    """Test that custom class types are correctly passed as result_type."""
    class CustomType:
        pass
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    update_method = DummyUpdateMethodMultiParam(ret_type=CustomType)
    codeflash_output = handle.get_update_handle_for(update_method, "update"); update_handle = codeflash_output # 500ns -> 500ns (0.000% faster)

def test_handle_with_none_return_type():
    """Test that None as return type is handled correctly."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    update_method = DummyUpdateMethodMultiParam(ret_type=None)
    codeflash_output = handle.get_update_handle_for(update_method, "update"); update_handle = codeflash_output # 542ns -> 458ns (18.3% faster)

def test_handle_with_long_string_ids():
    """Test that very long string IDs are handled correctly."""
    client = DummyClient()
    long_id = "x" * 1000
    handle = WorkflowHandle(client, long_id, run_id=long_id)
    update_method = DummyUpdateMethodMultiParam(ret_type=int)
    codeflash_output = handle.get_update_handle_for(update_method, long_id); update_handle = codeflash_output # 500ns -> 459ns (8.93% faster)

# -------- LARGE SCALE TEST CASES --------

def test_large_number_of_handles():
    """Test creating a large number of handles for scalability."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    update_method = DummyUpdateMethodMultiParam(ret_type=int)
    num_handles = 1000
    handles = []
    for i in range(num_handles):
        update_id = f"update_{i}"
        codeflash_output = handle.get_update_handle_for(update_method, update_id); update_handle = codeflash_output # 246μs -> 212μs (16.1% faster)
        handles.append(update_handle)

def test_large_variety_of_return_types():
    """Test handles with many different return types."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    types = [int, str, float, dict, list, set, tuple, bool, type(None)]
    # Add 991 unique dummy types to reach 1000
    for i in range(991):
        types.append(type(f"DummyType_{i}", (), {}))
    for idx, ret_type in enumerate(types):
        update_method = DummyUpdateMethodMultiParam(ret_type=ret_type)
        update_id = f"update_{idx}"
        codeflash_output = handle.get_update_handle_for(update_method, update_id); update_handle = codeflash_output # 245μs -> 210μs (16.9% faster)

def test_large_handle_with_large_ids():
    """Test handles with large data in IDs."""
    client = DummyClient()
    large_workflow_id = "workflow_" + "A" * 500
    large_run_id = "run_" + "B" * 500
    handle = WorkflowHandle(client, large_workflow_id, run_id=large_run_id)
    large_update_id = "update_" + "C" * 500
    update_method = DummyUpdateMethodMultiParam(ret_type=float)
    codeflash_output = handle.get_update_handle_for(update_method, large_update_id); update_handle = codeflash_output # 542ns -> 458ns (18.3% faster)

# -------- DETERMINISM TEST CASE --------

def test_determinism_with_same_inputs():
    """Test that calling with same inputs always produces same outputs."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    update_method = DummyUpdateMethodMultiParam(ret_type=int)
    update_id = "update"
    codeflash_output = handle.get_update_handle_for(update_method, update_id); handle1 = codeflash_output # 500ns -> 459ns (8.93% faster)
    codeflash_output = handle.get_update_handle_for(update_method, update_id); handle2 = codeflash_output # 292ns -> 250ns (16.8% faster)

# -------- ERROR HANDLING TEST CASES --------

def test_missing_defn_ret_type_raises():
    """Test that missing _defn.ret_type raises AttributeError."""
    client = DummyClient()
    handle = WorkflowHandle(client, "workflow", run_id="run")
    class BadUpdateMethod:
        pass
    bad_update = BadUpdateMethod()
    with pytest.raises(AttributeError):
        handle.get_update_handle_for(bad_update, "update") # 709ns -> 750ns (5.47% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from temporalio.client import WorkflowHandle

To edit these changes git checkout codeflash/optimize-WorkflowHandle.get_update_handle_for-mgh9xylx and push.