Made workflow_upsert_memo operate on info.raw_memo. Linter fixes.

Author: maciejdudko

temporalio/worker/_workflow_instance.py

@@ -218,9 +218,6 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._current_history_length = 0
         self._current_history_size = 0
         self._continue_as_new_suggested = False
-        self._raw_memo: Optional[
-            MutableMapping[str, temporalio.api.common.v1.Payload]
-        ] = dict(self._info.raw_memo.items())
         # Lazily loaded
         self._untyped_converted_memo: Optional[Mapping[str, Any]] = None
         # Handles which are ready to run on the next event loop iteration
@@ -1072,19 +1069,22 @@ def workflow_memo(self) -> Mapping[str, Any]:
         if self._untyped_converted_memo is None:
             self._untyped_converted_memo = {
                 k: self._payload_converter.from_payload(v)
-                for k, v in self._raw_memo.items()
+                for k, v in self._info.raw_memo.items()
             }
         return self._untyped_converted_memo
 
     def workflow_memo_value(
         self, key: str, default: Any, *, type_hint: Optional[Type]
     ) -> Any:
-        payload = self._raw_memo.get(key)
+        payload = self._info.raw_memo.get(key)
         if not payload:
             if default is temporalio.common._arg_unset:
                 raise KeyError(f"Memo does not have a value for key {key}")
             return default
-        return self._payload_converter.from_payload(payload, type_hint)
+        return self._payload_converter.from_payload(
+            payload,
+            type_hint,  # type: ignore[arg-type]
+        )
 
     def workflow_upsert_memo(self, updates: Mapping[str, Any]) -> None:
         # Converting before creating a command so that we don't leave a partial command in case of conversion failure.
@@ -1103,13 +1103,19 @@ def workflow_upsert_memo(self, updates: Mapping[str, Any]) -> None:
         command = self._add_command()
         fields = command.modify_workflow_properties.upserted_memo.fields
 
+        # Updating memo inside info by downcasting to mutable mapping.
+        mut_raw_memo = cast(
+            MutableMapping[str, temporalio.api.common.v1.Payload],
+            self._info.raw_memo,
+        )
+
         for k, v in update_payloads.items():
             fields[k].CopyFrom(v)
-            self._raw_memo[k] = v
+            mut_raw_memo[k] = v
 
         for k in removals:
             fields.get_or_create(k)
-            self._raw_memo.pop(k, None)
+            mut_raw_memo.pop(k, None)
 
         # Clearing cached value, will be regenerated on next workflow_memo() call.
         self._untyped_converted_memo = None