fix(sessions): annotate result_processor for mypy-diff; pyink-format test file

- Add Any type hints to DynamicPickleType.result_processor and its inner process() to clear mypy-diff [no-untyped-def].
- Reformat test_safe_unpickle.py to 2-space pyink style.

Author: daridor9

src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py

@@ -30,7 +30,6 @@
 from google.adk.sessions import _session_util
 from google.adk.sessions.migration import _schema_check_utils
 from google.adk.sessions.schemas import v1
-from google.adk.sessions.schemas._safe_unpickle import safe_loads as _safe_pickle_loads
 from google.genai import types
 import sqlalchemy
 from sqlalchemy import create_engine

src/google/adk/sessions/schemas/v0.py

@@ -111,11 +111,11 @@ def process_bind_param(self, value, dialect):
         return pickle.dumps(value)
     return value
 
-  def result_processor(self, dialect, coltype):
+  def result_processor(self, dialect: Any, coltype: Any) -> Any:
     if dialect.name in ("mysql", "spanner+spanner"):
       return super().result_processor(dialect, coltype)
 
-    def process(value):
+    def process(value: Any) -> Any:
       if value is None:
         return None
       if isinstance(value, memoryview):

tests/unittests/sessions/test_safe_unpickle.py

@@ -27,162 +27,164 @@
 
 
 def _make_global_payload(module: str, func: str, *args: str) -> bytes:
-    """Craft a pickle stream that calls module.func(*args)."""
-    buf = io.BytesIO()
-    buf.write(pickle.PROTO + struct.pack("B", 2))
-    buf.write(b"c" + f"{module}\n{func}\n".encode())
-    buf.write(b"(")
-    for arg in args:
-        encoded = arg.encode("utf-8")
-        buf.write(
-            pickle.SHORT_BINUNICODE + struct.pack("<B", len(encoded)) + encoded
-        )
-    buf.write(b"t")
-    buf.write(b"R")
-    buf.write(b".")
-    return buf.getvalue()
+  """Craft a pickle stream that calls module.func(*args)."""
+  buf = io.BytesIO()
+  buf.write(pickle.PROTO + struct.pack("B", 2))
+  buf.write(b"c" + f"{module}\n{func}\n".encode())
+  buf.write(b"(")
+  for arg in args:
+    encoded = arg.encode("utf-8")
+    buf.write(
+        pickle.SHORT_BINUNICODE + struct.pack("<B", len(encoded)) + encoded
+    )
+  buf.write(b"t")
+  buf.write(b"R")
+  buf.write(b".")
+  return buf.getvalue()
 
 
 class TestBlockedPayloads(unittest.TestCase):
-    """Malicious pickle payloads must be blocked."""
+  """Malicious pickle payloads must be blocked."""
 
-    def test_os_system(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(_make_global_payload("os", "system", "echo pwned"))
+  def test_os_system(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(_make_global_payload("os", "system", "echo pwned"))
 
-    def test_subprocess_popen(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(_make_global_payload("subprocess", "Popen", "id"))
+  def test_subprocess_popen(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(_make_global_payload("subprocess", "Popen", "id"))
 
-    def test_builtins_import(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(_make_global_payload("builtins", "__import__", "os"))
+  def test_builtins_import(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(_make_global_payload("builtins", "__import__", "os"))
 
-    def test_posix_system(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(_make_global_payload("posix", "system", "whoami"))
+  def test_posix_system(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(_make_global_payload("posix", "system", "whoami"))
 
-    def test_nt_system(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(_make_global_payload("nt", "system", "whoami"))
+  def test_nt_system(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(_make_global_payload("nt", "system", "whoami"))
 
-    def test_builtins_eval(self):
-        with self.assertRaises(pickle.UnpicklingError):
-            safe_loads(
-                _make_global_payload(
-                    "builtins", "eval", "__import__('os').system('id')"
-                )
-            )
+  def test_builtins_eval(self):
+    with self.assertRaises(pickle.UnpicklingError):
+      safe_loads(
+          _make_global_payload(
+              "builtins", "eval", "__import__('os').system('id')"
+          )
+      )
 
 
 class TestEventActionsRoundTrip(unittest.TestCase):
-    """Legitimate EventActions data must survive pickle -> safe_loads."""
-
-    def _round_trip(self, obj):
-        return safe_loads(pickle.dumps(obj))
-
-    def test_string_values(self):
-        original = {"state_delta": {"key": "value"}, "artifact_delta": {}}
-        self.assertEqual(self._round_trip(original), original)
-
-    def test_nested_dict(self):
-        original = {
-            "state_delta": {
-                "user_prefs": {"theme": "dark", "lang": "en"},
-                "counter": 42,
-            },
-            "artifact_delta": {"files": ["a.txt", "b.txt"]},
-        }
-        self.assertEqual(self._round_trip(original), original)
-
-    def test_none_and_bool(self):
-        original = {
-            "skip_summarization": True,
-            "requested_auth_configs": None,
-            "escalate": False,
-        }
-        self.assertEqual(self._round_trip(original), original)
-
-    def test_empty_dict(self):
-        self.assertEqual(self._round_trip({}), {})
+  """Legitimate EventActions data must survive pickle -> safe_loads."""
+
+  def _round_trip(self, obj):
+    return safe_loads(pickle.dumps(obj))
+
+  def test_string_values(self):
+    original = {"state_delta": {"key": "value"}, "artifact_delta": {}}
+    self.assertEqual(self._round_trip(original), original)
+
+  def test_nested_dict(self):
+    original = {
+        "state_delta": {
+            "user_prefs": {"theme": "dark", "lang": "en"},
+            "counter": 42,
+        },
+        "artifact_delta": {"files": ["a.txt", "b.txt"]},
+    }
+    self.assertEqual(self._round_trip(original), original)
+
+  def test_none_and_bool(self):
+    original = {
+        "skip_summarization": True,
+        "requested_auth_configs": None,
+        "escalate": False,
+    }
+    self.assertEqual(self._round_trip(original), original)
+
+  def test_empty_dict(self):
+    self.assertEqual(self._round_trip({}), {})
+
 
 class TestRealEventActionsRoundTrip(unittest.TestCase):
-    """Smoke test: real EventActions instances survive pickle -> safe_loads."""
-
-    def _round_trip(self, obj):
-        return safe_loads(pickle.dumps(obj))
-
-    def test_minimal_event_actions(self):
-        original = EventActions()
-        result = self._round_trip(original)
-        self.assertIsInstance(result, EventActions)
-        self.assertEqual(result.state_delta, {})
-        self.assertEqual(result.artifact_delta, {})
-
-    def test_event_actions_with_state_delta(self):
-        original = EventActions(
-            state_delta={"user_name": "alice", "turn_count": 3, "active": True},
-            artifact_delta={"report.pdf": 2},
-        )
-        result = self._round_trip(original)
-        self.assertIsInstance(result, EventActions)
-        self.assertEqual(result.state_delta, original.state_delta)
-        self.assertEqual(result.artifact_delta, original.artifact_delta)
-
-    def test_event_actions_with_transfer_and_escalate(self):
-        original = EventActions(
-            transfer_to_agent="specialist_agent",
-            escalate=True,
-            skip_summarization=True,
-        )
-        result = self._round_trip(original)
-        self.assertIsInstance(result, EventActions)
-        self.assertEqual(result.transfer_to_agent, "specialist_agent")
-        self.assertTrue(result.escalate)
-        self.assertTrue(result.skip_summarization)
-
-    def test_event_actions_with_complex_state_values(self):
-        original = EventActions(
-            state_delta={
-                "nested": {"a": [1, 2, 3], "b": None},
-                "count": 42,
-                "tags": ["ml", "security"],
-            },
-        )
-        result = self._round_trip(original)
-        self.assertIsInstance(result, EventActions)
-        self.assertEqual(result.state_delta["nested"]["a"], [1, 2, 3])
-        self.assertIsNone(result.state_delta["nested"]["b"])
+  """Smoke test: real EventActions instances survive pickle -> safe_loads."""
+
+  def _round_trip(self, obj):
+    return safe_loads(pickle.dumps(obj))
+
+  def test_minimal_event_actions(self):
+    original = EventActions()
+    result = self._round_trip(original)
+    self.assertIsInstance(result, EventActions)
+    self.assertEqual(result.state_delta, {})
+    self.assertEqual(result.artifact_delta, {})
+
+  def test_event_actions_with_state_delta(self):
+    original = EventActions(
+        state_delta={"user_name": "alice", "turn_count": 3, "active": True},
+        artifact_delta={"report.pdf": 2},
+    )
+    result = self._round_trip(original)
+    self.assertIsInstance(result, EventActions)
+    self.assertEqual(result.state_delta, original.state_delta)
+    self.assertEqual(result.artifact_delta, original.artifact_delta)
+
+  def test_event_actions_with_transfer_and_escalate(self):
+    original = EventActions(
+        transfer_to_agent="specialist_agent",
+        escalate=True,
+        skip_summarization=True,
+    )
+    result = self._round_trip(original)
+    self.assertIsInstance(result, EventActions)
+    self.assertEqual(result.transfer_to_agent, "specialist_agent")
+    self.assertTrue(result.escalate)
+    self.assertTrue(result.skip_summarization)
+
+  def test_event_actions_with_complex_state_values(self):
+    original = EventActions(
+        state_delta={
+            "nested": {"a": [1, 2, 3], "b": None},
+            "count": 42,
+            "tags": ["ml", "security"],
+        },
+    )
+    result = self._round_trip(original)
+    self.assertIsInstance(result, EventActions)
+    self.assertEqual(result.state_delta["nested"]["a"], [1, 2, 3])
+    self.assertIsNone(result.state_delta["nested"]["b"])
 
 
 class TestEnvVarFallback(unittest.TestCase):
-    """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler."""
-
-    _ENV_KEY = "ADK_ALLOW_UNSAFE_V0_PICKLE"
-    _PAYLOAD = _make_global_payload("collections", "Counter")
-
-    def test_blocked_without_env_var(self):
-        old = os.environ.pop(self._ENV_KEY, None)
-        try:
-            with self.assertRaises(pickle.UnpicklingError):
-                safe_loads(self._PAYLOAD)
-        finally:
-            if old is not None:
-                os.environ[self._ENV_KEY] = old
-
-    def test_allowed_with_env_var(self):
-        old = os.environ.get(self._ENV_KEY)
-        try:
-            os.environ[self._ENV_KEY] = "1"
-            from collections import Counter
-            result = safe_loads(self._PAYLOAD)
-            self.assertIsInstance(result, Counter)
-        finally:
-            if old is None:
-                os.environ.pop(self._ENV_KEY, None)
-            else:
-                os.environ[self._ENV_KEY] = old
+  """ADK_ALLOW_UNSAFE_V0_PICKLE=1 must bypass RestrictedUnpickler."""
+
+  _ENV_KEY = "ADK_ALLOW_UNSAFE_V0_PICKLE"
+  _PAYLOAD = _make_global_payload("collections", "Counter")
+
+  def test_blocked_without_env_var(self):
+    old = os.environ.pop(self._ENV_KEY, None)
+    try:
+      with self.assertRaises(pickle.UnpicklingError):
+        safe_loads(self._PAYLOAD)
+    finally:
+      if old is not None:
+        os.environ[self._ENV_KEY] = old
+
+  def test_allowed_with_env_var(self):
+    old = os.environ.get(self._ENV_KEY)
+    try:
+      os.environ[self._ENV_KEY] = "1"
+      from collections import Counter
+
+      result = safe_loads(self._PAYLOAD)
+      self.assertIsInstance(result, Counter)
+    finally:
+      if old is None:
+        os.environ.pop(self._ENV_KEY, None)
+      else:
+        os.environ[self._ENV_KEY] = old
 
 
 if __name__ == "__main__":
-    unittest.main()
+  unittest.main()