PR Feedback

Author: andystaples

azure/durable_functions/models/DurableEntityContext.py

@@ -108,13 +108,9 @@ def from_json(cls, json_str: str) -> Tuple['DurableEntityContext', List[Dict[str
         json_dict["key"] = json_dict["self"]["key"]
         json_dict.pop("self")
 
+        # Keep the raw serialized state (a JSON string) so get_state() can
+        # deserialize lazily with an expected_type supplied by the user.
         serialized_state = json_dict["state"]
-        if serialized_state is not None:
-            # Keep the raw serialized form so get_state() can deserialize
-            # lazily with an expected_type supplied by the user.
-            json_dict["state"] = serialized_state
-        else:
-            json_dict["state"] = None
 
         batch = json_dict.pop("batch")
         ctx = cls(**json_dict)
@@ -134,6 +130,10 @@ def set_state(self, state: Any) -> None:
 
         # should only serialize the state at the end of the batch
         self._state = state
+        # The new state is a live Python value, not the raw JSON string
+        # loaded from the payload. Clear the raw flag so a subsequent
+        # get_state() in the same batch does not try to re-decode it.
+        self._state_is_raw = False
 
     def get_state(self, initializer: Optional[Callable[[], Any]] = None,
                   expected_type: Optional[type] = None) -> Any:
@@ -145,7 +145,11 @@ def get_state(self, initializer: Optional[Callable[[], Any]] = None,
             A 0-argument function to provide an initial state. Defaults to None.
         expected_type: Optional[type]
             The type to decode the state as. When set, the codec uses
-            this type directly without consulting ``sys.modules``.
+            this type directly without consulting ``sys.modules``. Note that
+            the persisted state is decoded lazily on the **first** get_state
+            call within a batch; an ``expected_type`` supplied on a later
+            call (after the state has already been decoded or replaced via
+            set_state) has no effect.
 
         Returns
         -------
@@ -199,6 +203,7 @@ def destruct_on_exit(self) -> None:
         """Delete this entity after the operation completes."""
         self._exists = False
         self._state = None
+        self._state_is_raw = False
 
 
 def from_json_util(json_str: str, expected_type: Optional[type] = None) -> Any:

azure/durable_functions/models/OrchestratorState.py

@@ -1,10 +1,9 @@
-import json
 from typing import List, Any, Dict, Optional
 
 from azure.durable_functions.models.ReplaySchema import ReplaySchema
 
 from .utils.json_utils import add_attrib
-from .utils.df_serialization import _get_serialize_default
+from .utils.df_serialization import df_dumps
 from azure.durable_functions.models.actions.Action import Action
 
 
@@ -114,4 +113,4 @@ def to_json_string(self) -> str:
             The instance of the object in json string format
         """
         json_dict = self.to_json()
-        return json.dumps(json_dict, default=_get_serialize_default())
+        return df_dumps(json_dict)

azure/durable_functions/models/utils/type_discovery.py

@@ -12,8 +12,10 @@
 
 from __future__ import annotations
 
+import functools
 import inspect
 import logging
+import typing
 from typing import Any, Callable, Optional
 
 logger = logging.getLogger(__name__)
@@ -33,12 +35,35 @@ def _unwrap_function_builder(name_or_callable: Any) -> Optional[Callable]:
     return None
 
 
+@functools.lru_cache(maxsize=None)
 def _return_annotation(fn: Callable) -> Optional[type]:
+    """Resolve *fn*'s return annotation to a concrete ``type``, or ``None``.
+
+    ``typing.get_type_hints`` is tried first so that string annotations
+    (``from __future__ import annotations`` / PEP 563) are resolved to the
+    real object. Results are memoized per function because this runs on
+    every ``call_activity`` / ``call_sub_orchestrator`` (including replay).
+
+    Limitation: generic aliases such as ``list[Order]`` or
+    ``Optional[Order]`` are not concrete ``type`` objects, so they resolve
+    to ``None`` and the caller falls back to module-only resolution.
+    """
+    ann: Any = inspect.Signature.empty
     try:
-        sig = inspect.signature(fn)
-    except (TypeError, ValueError):
-        return None
-    ann = sig.return_annotation
+        hints = typing.get_type_hints(fn)
+    except Exception:
+        hints = None
+    if hints is not None and "return" in hints:
+        ann = hints["return"]
+    else:
+        # get_type_hints couldn't resolve (e.g. forward ref it can't see);
+        # fall back to the raw signature annotation.
+        try:
+            sig = inspect.signature(fn)
+        except (TypeError, ValueError):
+            return None
+        ann = sig.return_annotation
+
     if ann is inspect.Signature.empty:
         return None
     return ann if isinstance(ann, type) else None

eng/templates/build.yml

@@ -64,4 +64,4 @@ jobs:
         - script: |
             pip install pytest pytest-azurepipelines
             pytest --ignore=samples-v2
-          displayName: 'pytest'
+          displayName: 'pytest'

tests/orchestrator/test_entity.py

@@ -94,6 +94,18 @@ def counter_entity_function_raises_exception(context):
 def counter_entity_function_raises_exception_with_pystein(context):
     raise Exception("boom!")
 
+def set_then_get_entity(context):
+    """Entity that sets state (without first reading it) in one operation and
+    reads it in a later operation. Used to exercise set-then-get across a
+    batch when the entity already has persisted state.
+    """
+    operation = context.operation_name
+    if operation == "set":
+        context.set_state(10)
+        context.set_result("set")
+    elif operation == "get":
+        context.set_result(context.get_state(lambda: 0))
+
 def test_entity_raises_exception():
     # Create input batch
     batch = []
@@ -163,6 +175,40 @@ def test_entity_signal_then_call():
     #assert_valid_schema(result)
     assert_entity_state_equals(expected, result)
 
+def test_entity_set_then_get_with_preexisting_raw_state():
+    """Regression test: an entity that already has persisted state must be
+    able to set_state in one operation and get_state in a later operation
+    within the same batch.
+
+    ``from_json`` keeps the persisted state in its raw (undecoded) form and
+    marks it as raw so the first ``get_state`` can decode it lazily with a
+    user-supplied ``expected_type``. ``set_state`` replaces that raw value
+    with a live Python value, so it must clear the raw flag -- otherwise a
+    later ``get_state`` would try to re-decode an already-live value and the
+    operation would fail.
+    """
+    # Pre-existing persisted state (single-encoded JSON string) is what makes
+    # from_json mark the loaded state as raw.
+    batch = []
+    add_to_batch(batch, name="set")
+    add_to_batch(batch, name="get")
+    context_builder = EntityContextBuilder(batch=batch, state=json.dumps(5))
+
+    # Run the entity, get observed result
+    result = get_entity_state_result(
+        context_builder,
+        set_then_get_entity,
+        )
+
+    # Both operations should succeed; the "get" must observe the value set by
+    # the earlier "set" (10), not crash trying to re-decode it.
+    expected_state = entity_base_expected_state()
+    apply_operation(expected_state, result="set", state=10)
+    apply_operation(expected_state, result=10, state=10)
+    expected = expected_state.to_json()
+
+    assert_entity_state_equals(expected, result)
+
 def test_entity_signal_then_call_with_pystein():
     """Tests that a simple counter entity outputs the correct value
     after a sequence of operations. Mostly just a sanity check.

tests/orchestrator/test_external_event.py

@@ -85,4 +85,4 @@ def test_external_event_with_expected_type():
         context_builder, generator_function_with_expected_type)
 
     assert result["isDone"] is True
-    assert result["output"] == "hello"
+    assert result["output"] == "hello"