Add conversion functions as constructor parameters

sokoliva · sokoliva · commit 97b82e241d8d · 2026-03-16T16:09:09.000Z
diff --git a/src/a2a/server/tasks/database_push_notification_config_store.py b/src/a2a/server/tasks/database_push_notification_config_store.py
@@ -1,5 +1,4 @@
 # ruff: noqa: PLC0415
-import inspect
 import logging
 
 from typing import TYPE_CHECKING
@@ -25,8 +24,7 @@
         "or 'pip install a2a-sdk[sql]'"
     ) from e
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
+from collections.abc import Callable
 
 from a2a.compat.v0_3.conversions import (
     compat_push_notification_config_model_to_core,
@@ -62,17 +60,34 @@ class DatabasePushNotificationConfigStore(PushNotificationConfigStore):
     config_model: type[PushNotificationConfigModel]
     _fernet: 'Fernet | None'
     owner_resolver: OwnerResolver
+    core_to_model_conversion: (
+        Callable[
+            [str, TaskPushNotificationConfig, str, 'Fernet | None'],
+            PushNotificationConfigModel,
+        ]
+        | None
+    )
+    model_to_core_conversion: (
+        Callable[[PushNotificationConfigModel], TaskPushNotificationConfig]
+        | None
+    )
 
-    core_to_model_conversion: 'Callable[[str, TaskPushNotificationConfig, str, Fernet | None], PushNotificationConfigModel] | None' = None
-    model_to_core_conversion: 'Callable[[PushNotificationConfigModel], TaskPushNotificationConfig] | None' = None
-
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         engine: AsyncEngine,
         create_table: bool = True,
         table_name: str = 'push_notification_configs',
         encryption_key: str | bytes | None = None,
         owner_resolver: OwnerResolver = resolve_user_scope,
+        core_to_model_conversion: Callable[
+            [str, TaskPushNotificationConfig, str, 'Fernet | None'],
+            PushNotificationConfigModel,
+        ]
+        | None = None,
+        model_to_core_conversion: Callable[
+            [PushNotificationConfigModel], TaskPushNotificationConfig
+        ]
+        | None = None,
     ) -> None:
         """Initializes the DatabasePushNotificationConfigStore.
 
@@ -84,6 +99,8 @@ def __init__(
                 If provided, `config_data` will be encrypted in the database.
                 The key must be a URL-safe base64-encoded 32-byte key.
             owner_resolver: Function to resolve the owner from the context.
+            core_to_model_conversion: Optional function to convert a TaskPushNotificationConfig to a TaskPushNotificationConfigModel.
+            model_to_core_conversion: Optional function to convert a TaskPushNotificationConfigModel to a TaskPushNotificationConfig.
         """
         logger.debug(
             'Initializing DatabasePushNotificationConfigStore with existing engine, table: %s',
@@ -102,6 +119,8 @@ def __init__(
             else create_push_notification_config_model(table_name)
         )
         self._fernet = None
+        self.core_to_model_conversion = core_to_model_conversion
+        self.model_to_core_conversion = model_to_core_conversion
 
         if encryption_key:
             try:
@@ -156,12 +175,10 @@ def _to_orm(
 
         The config data is serialized to JSON bytes, and encrypted if a key is configured.
         """
-        if conversion := self.core_to_model_conversion:
-            # If it's a bound method of this instance, call the underlying function
-            # to avoid passing 'self' twice.
-            if inspect.ismethod(conversion):
-                return conversion.__func__(task_id, config, owner, self._fernet)
-            return conversion(task_id, config, owner, self._fernet)
+        if self.core_to_model_conversion:
+            return self.core_to_model_conversion(
+                task_id, config, owner, self._fernet
+            )
 
         json_payload = MessageToJson(config).encode('utf-8')
 
@@ -185,12 +202,8 @@ def _from_orm(
 
         Handles decryption if a key is configured, with a fallback to plain JSON.
         """
-        if conversion := self.model_to_core_conversion:
-            # If it's a bound method of this instance, call the underlying function
-            # to avoid passing 'self' twice.
-            if inspect.ismethod(conversion):
-                return conversion.__func__(model_instance)
-            return conversion(model_instance)
+        if self.model_to_core_conversion:
+            return self.model_to_core_conversion(model_instance)
 
         payload = model_instance.config_data
 
diff --git a/src/a2a/server/tasks/database_task_store.py b/src/a2a/server/tasks/database_task_store.py
@@ -1,8 +1,7 @@
-import inspect
 import logging
 
+from collections.abc import Callable
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING
 
 
 try:
@@ -22,13 +21,11 @@
         "'pip install a2a-sdk[sqlite]', "
         "or 'pip install a2a-sdk[sql]'"
     ) from e
-
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
 from google.protobuf.json_format import MessageToDict, ParseDict
 
-from a2a.compat.v0_3 import conversions
+from a2a.compat.v0_3.conversions import (
+    compat_task_model_to_core,
+)
 from a2a.server.context import ServerCallContext
 from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
@@ -55,16 +52,18 @@ class DatabaseTaskStore(TaskStore):
     _initialized: bool
     task_model: type[TaskModel]
     owner_resolver: OwnerResolver
+    core_to_model_conversion: Callable[[Task, str], TaskModel] | None = None
+    model_to_core_conversion: Callable[[TaskModel], Task] | None = None
 
-    core_to_model_conversion: 'Callable[[Task, str], TaskModel] | None' = None
-    model_to_core_conversion: 'Callable[[TaskModel], Task] | None' = None
-
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         engine: AsyncEngine,
         create_table: bool = True,
         table_name: str = 'tasks',
         owner_resolver: OwnerResolver = resolve_user_scope,
+        core_to_model_conversion: Callable[[Task, str], TaskModel]
+        | None = None,
+        model_to_core_conversion: Callable[[TaskModel], Task] | None = None,
     ) -> None:
         """Initializes the DatabaseTaskStore.
 
@@ -73,6 +72,8 @@ def __init__(
             create_table: If true, create tasks table on initialization.
             table_name: Name of the database table. Defaults to 'tasks'.
             owner_resolver: Function to resolve the owner from the context.
+            core_to_model_conversion: Optional function to convert a Task to a TaskModel.
+            model_to_core_conversion: Optional function to convert a TaskModel to a Task.
         """
         logger.debug(
             'Initializing DatabaseTaskStore with existing engine, table: %s',
@@ -85,6 +86,8 @@ def __init__(
         self.create_table = create_table
         self._initialized = False
         self.owner_resolver = owner_resolver
+        self.core_to_model_conversion = core_to_model_conversion
+        self.model_to_core_conversion = model_to_core_conversion
 
         self.task_model = (
             TaskModel
@@ -117,12 +120,8 @@ async def _ensure_initialized(self) -> None:
 
     def _to_orm(self, task: Task, owner: str) -> TaskModel:
         """Maps a Proto Task to a SQLAlchemy TaskModel instance."""
-        if conversion := self.core_to_model_conversion:
-            # If it's a bound method of this instance, call the underlying function
-            # to avoid passing 'self' twice.
-            if inspect.ismethod(conversion):
-                return conversion.__func__(task, owner)
-            return conversion(task, owner)
+        if self.core_to_model_conversion:
+            return self.core_to_model_conversion(task, owner)
 
         return self.task_model(
             id=task.id,
@@ -145,12 +144,8 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
 
     def _from_orm(self, task_model: TaskModel) -> Task:
         """Maps a SQLAlchemy TaskModel to a Proto Task instance."""
-        if conversion := self.model_to_core_conversion:
-            # If it's a bound method of this instance, call the underlying function
-            # to avoid passing 'self' twice.
-            if inspect.ismethod(conversion):
-                return conversion.__func__(task_model)
-            return conversion(task_model)
+        if self.model_to_core_conversion:
+            return self.model_to_core_conversion(task_model)
 
         if task_model.protocol_version == '1.0':
             task = Task(
@@ -172,7 +167,7 @@ def _from_orm(self, task_model: TaskModel) -> Task:
             return task
 
         # Legacy conversion
-        return conversions.compat_task_model_to_core(task_model)
+        return compat_task_model_to_core(task_model)
 
     async def save(
         self, task: Task, context: ServerCallContext | None = None
diff --git a/tests/server/tasks/test_database_push_notification_config_store.py b/tests/server/tasks/test_database_push_notification_config_store.py
@@ -788,7 +788,6 @@ async def test_get_0_3_push_notification_config_detailed(
 @pytest.mark.asyncio
 async def test_custom_conversion():
     engine = MagicMock()
-    store = DatabasePushNotificationConfigStore(engine=engine)
 
     # Custom callables
     mock_to_orm = MagicMock(
@@ -797,33 +796,26 @@ async def test_custom_conversion():
     mock_from_orm = MagicMock(
         return_value=TaskPushNotificationConfig(id='custom_config')
     )
+    store = DatabasePushNotificationConfigStore(
+        engine=engine,
+        core_to_model_conversion=mock_to_orm,
+        model_to_core_conversion=mock_from_orm,
+    )
 
-    DatabasePushNotificationConfigStore.core_to_model_conversion = mock_to_orm
-    DatabasePushNotificationConfigStore.model_to_core_conversion = mock_from_orm
-
-    try:
-        config = TaskPushNotificationConfig(id='orig')
-        model = store._to_orm('t1', config, 'owner')
-        assert model.config_id == 'c1'
-        mock_to_orm.assert_called_once_with('t1', config, 'owner', None)
+    config = TaskPushNotificationConfig(id='orig')
+    model = store._to_orm('t1', config, 'owner')
+    assert model.config_id == 'c1'
+    mock_to_orm.assert_called_once_with('t1', config, 'owner', None)
 
-        model_instance = PushNotificationConfigModel(
-            task_id='t1', config_id='c1'
-        )
-        loaded_config = store._from_orm(model_instance)
-        assert loaded_config.id == 'custom_config'
-        mock_from_orm.assert_called_once_with(model_instance)
-    finally:
-        # Reset class variables
-        DatabasePushNotificationConfigStore.core_to_model_conversion = None
-        DatabasePushNotificationConfigStore.model_to_core_conversion = None
+    model_instance = PushNotificationConfigModel(task_id='t1', config_id='c1')
+    loaded_config = store._from_orm(model_instance)
+    assert loaded_config.id == 'custom_config'
+    mock_from_orm.assert_called_once_with(model_instance)
 
 
-@pytest.mark.parametrize('assignment_type', ['class', 'instance'])
 @pytest.mark.asyncio
 async def test_core_to_0_3_model_conversion(
     db_store_parameterized: DatabasePushNotificationConfigStore,
-    assignment_type: str,
 ) -> None:
     """Test storing and retrieving push notification configs in v0.3 format using conversion utilities.
 
@@ -834,14 +826,9 @@ async def test_core_to_0_3_model_conversion(
     store = db_store_parameterized
 
     # Set the v0.3 persistence utilities
-    if assignment_type == 'class':
-        DatabasePushNotificationConfigStore.core_to_model_conversion = (
-            core_to_compat_push_notification_config_model
-        )
-    else:
-        store.core_to_model_conversion = (
-            core_to_compat_push_notification_config_model
-        )
+    store.core_to_model_conversion = (
+        core_to_compat_push_notification_config_model
+    )
 
     task_id = 'v03-persistence-task'
     config_id = 'c1'
@@ -852,6 +839,7 @@ async def test_core_to_0_3_model_conversion(
     )
     # 1. Save the config (will use core_to_compat_push_notification_config_model)
     await store.set_info(task_id, original_config, MINIMAL_CALL_CONTEXT)
+
     # 2. Verify it's stored in v0.3 format directly in DB
     async with store.async_session_maker() as session:
         db_model = await session.get(store.config_model, (task_id, config_id))
@@ -868,16 +856,15 @@ async def test_core_to_0_3_model_conversion(
         assert data['id'] == 'c1'
         assert data['token'] == 'legacy-token'
         assert 'taskId' not in data
+
     # 3. Retrieve the config (will use compat_push_notification_config_model_to_core)
     retrieved_configs = await store.get_info(task_id, MINIMAL_CALL_CONTEXT)
     assert len(retrieved_configs) == 1
     retrieved = retrieved_configs[0]
     assert retrieved.id == original_config.id
     assert retrieved.url == original_config.url
     assert retrieved.token == original_config.token
+
     # Reset conversion attributes
-    if assignment_type == 'class':
-        DatabasePushNotificationConfigStore.core_to_model_conversion = None
-    else:
-        store.core_to_model_conversion = None
+    store.core_to_model_conversion = None
     await store.delete_info(task_id, MINIMAL_CALL_CONTEXT)
diff --git a/tests/server/tasks/test_database_task_store.py b/tests/server/tasks/test_database_task_store.py
