Address Romain review feedback

- Add type annotations to ArtifactSerializer method signatures for better
  editor/IDE support.
- Remove unused ``context`` parameter from ``deserialize`` (call sites
  never passed a meaningful value).
- Promote ``STORAGE`` / ``WIRE`` to a ``SerializationFormat`` enum;
  subclass ``str`` so the existing string-literal comparisons keep working.
- Reword the ``_serializers_override`` comment and move the property
  rationale next to the property definition.
- Extend the multi-blob error message with "at this time. If you have a
  need for multi blob serializers, please reach out to the Metaflow
  team.".
- Include ``serializer_info`` in the "No deserializer claimed artifact"
  error message.

Author: Saeid Barati

metaflow/datastore/artifacts/serializer.py

@@ -1,12 +1,32 @@
 import inspect
 from abc import ABCMeta, abstractmethod
 from collections import namedtuple
+from enum import Enum
+from typing import Any, List, Optional, Tuple, Type, Union
 
 
-# Serialization formats. STORAGE produces (blobs, metadata) for the datastore;
-# WIRE produces a str for CLI args, protobuf payloads, and cross-process IPC.
-STORAGE = "storage"
-WIRE = "wire"
+class SerializationFormat(str, Enum):
+    """
+    Representation a serializer produces or consumes.
+
+    - ``STORAGE`` yields ``(List[SerializedBlob], SerializationMetadata)`` for
+      the datastore save path.
+    - ``WIRE`` yields a ``str`` for CLI args, protobuf payloads, and
+      cross-process IPC.
+
+    Subclassing ``str`` keeps ``SerializationFormat.STORAGE == "storage"``
+    True, so existing call sites that compare against the string literal
+    keep working without a migration.
+    """
+
+    STORAGE = "storage"
+    WIRE = "wire"
+
+
+# Module-level aliases kept so call sites can write ``format=STORAGE`` without
+# importing the enum itself.
+STORAGE = SerializationFormat.STORAGE
+WIRE = SerializationFormat.WIRE
 
 
 SerializationMetadata = namedtuple(
@@ -30,7 +50,11 @@ class SerializedBlob(object):
         If None, auto-detected from value type: str -> reference, bytes -> new data.
     """
 
-    def __init__(self, value, is_reference=None):
+    def __init__(
+        self,
+        value: Union[str, bytes],
+        is_reference: Optional[bool] = None,
+    ):
         if not isinstance(value, (str, bytes)):
             raise TypeError(
                 "SerializedBlob value must be str or bytes, got %s" % type(value).__name__
@@ -42,7 +66,7 @@ def __init__(self, value, is_reference=None):
             self.is_reference = is_reference
 
     @property
-    def needs_save(self):
+    def needs_save(self) -> bool:
         """True if this blob contains new bytes that need to be stored."""
         return not self.is_reference
 
@@ -69,7 +93,7 @@ def __init__(cls, name, bases, namespace):
         SerializerStore._ordered_cache = None
 
     @staticmethod
-    def get_ordered_serializers():
+    def get_ordered_serializers() -> List[Type["ArtifactSerializer"]]:
         """
         Return serializer classes sorted by (PRIORITY, registration_order).
 
@@ -127,12 +151,12 @@ class ArtifactSerializer(object, metaclass=SerializerStore):
         PickleSerializer uses 9999 as the universal fallback.
     """
 
-    TYPE = None
-    PRIORITY = 100
+    TYPE: Optional[str] = None
+    PRIORITY: int = 100
 
     @classmethod
     @abstractmethod
-    def can_serialize(cls, obj):
+    def can_serialize(cls, obj: Any) -> bool:
         """
         Return True if this serializer can handle the given object.
 
@@ -149,7 +173,7 @@ def can_serialize(cls, obj):
 
     @classmethod
     @abstractmethod
-    def can_deserialize(cls, metadata):
+    def can_deserialize(cls, metadata: SerializationMetadata) -> bool:
         """
         Return True if this serializer can deserialize given the metadata.
 
@@ -166,7 +190,11 @@ def can_deserialize(cls, metadata):
 
     @classmethod
     @abstractmethod
-    def serialize(cls, obj, format=STORAGE):
+    def serialize(
+        cls,
+        obj: Any,
+        format: SerializationFormat = STORAGE,
+    ) -> Union[Tuple[List[SerializedBlob], SerializationMetadata], str]:
         """
         Serialize obj. Must be side-effect-free: this method may be invoked
         multiple times (caching, retries, parallel dispatch) and must not
@@ -178,7 +206,7 @@ def serialize(cls, obj, format=STORAGE):
         ----------
         obj : Any
             The Python object to serialize.
-        format : str
+        format : SerializationFormat
             Either ``STORAGE`` (default) or ``WIRE``.
             - ``STORAGE`` returns a tuple ``(List[SerializedBlob], SerializationMetadata)``
               for persisting through the datastore.
@@ -194,7 +222,12 @@ def serialize(cls, obj, format=STORAGE):
 
     @classmethod
     @abstractmethod
-    def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
+    def deserialize(
+        cls,
+        data: Union[List[bytes], str],
+        metadata: Optional[SerializationMetadata] = None,
+        format: SerializationFormat = STORAGE,
+    ) -> Any:
         """
         Deserialize back to a Python object.
 
@@ -205,9 +238,7 @@ def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
         metadata : SerializationMetadata, optional
             Metadata stored alongside the artifact. Required for STORAGE,
             ignored for WIRE.
-        context : Any, optional
-            Optional context for deserialization (e.g., task vs client loading).
-        format : str
+        format : SerializationFormat
             Either ``STORAGE`` (default) or ``WIRE``.
 
         Returns

metaflow/datastore/task_datastore.py

@@ -117,12 +117,9 @@ def __init__(
         self._parent = flow_datastore
         self._persist = persist
 
-        # ``_serializers`` is a property that dispatches through
-        # ``SerializerStore.get_ordered_serializers()`` on each access. The
-        # lookup is cheap (cached inside the store) and picks up serializers
-        # registered via the lazy import hook after this instance was
-        # constructed — otherwise long-lived datastores (notebooks, client
-        # sessions) would silently miss any extension registered after init.
+        # Tests assign ``self._serializers = [...]`` to pin the dispatch list
+        # for isolation. When set, the ``_serializers`` property returns this
+        # override instead of consulting the global registry.
         self._serializers_override = None
 
         self._is_done_set = False
@@ -203,6 +200,12 @@ def __init__(
 
     @property
     def _serializers(self):
+        # Dispatch through ``SerializerStore.get_ordered_serializers()`` on
+        # each access. The lookup is cheap (cached inside the store) and
+        # picks up serializers registered via the lazy import hook after
+        # this instance was constructed — otherwise long-lived datastores
+        # (notebooks, client sessions) would silently miss any extension
+        # registered after init.
         if self._serializers_override is not None:
             return self._serializers_override
         return SerializerStore.get_ordered_serializers()
@@ -402,7 +405,9 @@ def serialize_iter():
                     # load. Fail loudly until multi-blob support lands.
                     raise DataException(
                         "Serializer %s returned %d blobs for artifact '%s'; "
-                        "only single-blob serializers are supported."
+                        "only single-blob serializers are supported at this "
+                        "time. If you have a need for multi blob "
+                        "serializers, please reach out to the Metaflow team."
                         % (serializer.__name__, len(blobs), name)
                     )
                 artifact_names.append(name)
@@ -472,8 +477,14 @@ def load_artifacts(self, names):
                         % serializer_source
                     )
                 raise DataException(
-                    "No deserializer claimed artifact '%s' (encoding: %s)."
-                    "%s" % (name, metadata.encoding, source_hint)
+                    "No deserializer claimed artifact '%s' (encoding: %s, "
+                    "serializer_info: %r).%s"
+                    % (
+                        name,
+                        metadata.encoding,
+                        metadata.serializer_info,
+                        source_hint,
+                    )
                 )
             deserializers[name] = (deserializer, metadata)
             to_load[self._objects[name]].append(name)
@@ -486,9 +497,7 @@ def load_artifacts(self, names):
                 # Deserialize each time to have fully distinct objects (the user
                 # would not expect two artifacts with different names to actually
                 # be aliases of one another)
-                yield name, deserializer.deserialize(
-                    [blob], metadata, context=None
-                )
+                yield name, deserializer.deserialize([blob], metadata)
 
     @require_mode("r")
     def get_artifact_sizes(self, names):

metaflow/plugins/datastores/serializers/pickle_serializer.py

@@ -55,7 +55,7 @@ def serialize(cls, obj, format=STORAGE):
         )
 
     @classmethod
-    def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
+    def deserialize(cls, data, metadata=None, format=STORAGE):
         if format == WIRE:
             raise NotImplementedError(
                 "PickleSerializer does not support the WIRE format."

test/unit/test_artifact_serializer.py

@@ -51,7 +51,7 @@ def serialize(cls, obj):
         )
 
     @classmethod
-    def deserialize(cls, blobs, metadata, context):
+    def deserialize(cls, blobs, metadata):
         return blobs[0].decode("utf-8")
 
 
@@ -76,7 +76,7 @@ def serialize(cls, obj):
         )
 
     @classmethod
-    def deserialize(cls, blobs, metadata, context):
+    def deserialize(cls, blobs, metadata):
         return int(blobs[0].decode("utf-8"))
 
 
@@ -99,7 +99,7 @@ def serialize(cls, obj):
         raise NotImplementedError
 
     @classmethod
-    def deserialize(cls, blobs, metadata, context):
+    def deserialize(cls, blobs, metadata):
         raise NotImplementedError
 
 
@@ -143,7 +143,7 @@ def serialize(cls, obj):
                 raise NotImplementedError
 
             @classmethod
-            def deserialize(cls, blobs, metadata, context):
+            def deserialize(cls, blobs, metadata):
                 raise NotImplementedError
 
         assert SerializerStore._all_serializers["test_high"] is _ReplacementSerializer
@@ -286,7 +286,7 @@ def serialize(cls, obj, format=STORAGE):
         )
 
     @classmethod
-    def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
+    def deserialize(cls, data, metadata=None, format=STORAGE):
         if format == WIRE:
             return data
         return data[0].decode("utf-8")
@@ -295,6 +295,10 @@ def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
 def test_format_constants():
     assert STORAGE == "storage"
     assert WIRE == "wire"
+    # Enum members compare both by identity and against the underlying string.
+    assert STORAGE is STORAGE
+    assert STORAGE.value == "storage"
+    assert WIRE.value == "wire"
 
 
 def test_dual_format_storage_roundtrip():

test/unit/test_lazy_serializer_registry.py

@@ -200,7 +200,7 @@ def serialize(cls, obj, format="storage"):
             return [SerializedBlob(blob)], SerializationMetadata("x", 0, "x", {})
 
         @classmethod
-        def deserialize(cls, data, metadata=None, context=None, format="storage"):
+        def deserialize(cls, data, metadata=None, format="storage"):
             return None
 
     # Remove it so lazy-registry has to pull it in.

test/unit/test_pickle_serializer.py

@@ -165,7 +165,7 @@ def test_serialize_metadata_serializer_info_empty():
 def test_round_trip(obj):
     blobs, meta = PickleSerializer.serialize(obj)
     raw_blobs = [b.value for b in blobs]
-    result = PickleSerializer.deserialize(raw_blobs, meta, context=None)
+    result = PickleSerializer.deserialize(raw_blobs, meta)
     assert result == obj
 
 
@@ -181,6 +181,6 @@ def test_round_trip_custom_class():
     obj = _CustomObj(42)
     blobs, meta = PickleSerializer.serialize(obj)
     raw_blobs = [b.value for b in blobs]
-    result = PickleSerializer.deserialize(raw_blobs, meta, context=None)
+    result = PickleSerializer.deserialize(raw_blobs, meta)
     assert result == obj
     assert result.x == 42

test/unit/test_serializer_integration.py

@@ -141,7 +141,7 @@ def serialize(cls, obj):
             )
 
         @classmethod
-        def deserialize(cls, blobs, metadata, context):
+        def deserialize(cls, blobs, metadata):
             return json.loads(blobs[0].decode("utf-8"))
 
     # Explicitly set serializers: custom first, then pickle fallback.
@@ -240,7 +240,7 @@ def serialize(cls, obj, format="storage"):
             raise NotImplementedError
 
         @classmethod
-        def deserialize(cls, data, metadata=None, context=None, format="storage"):
+        def deserialize(cls, data, metadata=None, format="storage"):
             raise NotImplementedError
 
     try: