fix: address comments

Author: geoHeil

.basedpyright/baseline.json

@@ -9103,38 +9103,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnknownParameterType",
-                "range": {
-                    "startColumn": 60,
-                    "endColumn": 66,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportMissingParameterType",
-                "range": {
-                    "startColumn": 60,
-                    "endColumn": 66,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 36,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 36,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportImplicitStringConcatenation",
                 "range": {
@@ -9151,53 +9119,13 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportUnknownParameterType",
-                "range": {
-                    "startColumn": 60,
-                    "endColumn": 66,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportMissingParameterType",
-                "range": {
-                    "startColumn": 60,
-                    "endColumn": 66,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnusedCallResult",
                 "range": {
                     "startColumn": 8,
                     "endColumn": 43,
                     "lineCount": 1
                 }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 36,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 36,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 36,
-                    "endColumn": 42,
-                    "lineCount": 1
-                }
             }
         ],
         "./src/metaxy/models/field.py": [

src/metaxy/models/feature_spec.py

@@ -13,7 +13,6 @@
     Protocol,
     TypeAlias,
     TypeVar,
-    cast,
     overload,
     runtime_checkable,
 )
@@ -209,26 +208,6 @@ def table_name(self) -> str:
 )  # bound, should be used for generic
 
 
-def _coerce_metadata(value: Any) -> dict[str, JsonValue] | None:
-    if value is None:
-        return None
-    if not isinstance(value, Mapping):
-        raise ValueError("metadata must be a mapping")
-    try:
-        serialized = json.dumps(value)
-    except (TypeError, ValueError) as exc:
-        raise ValueError(
-            "metadata must be JSON-serializable. Found non-serializable value"
-        ) from exc
-    return cast(dict[str, JsonValue], json.loads(serialized))
-
-
-MetadataField = Annotated[
-    dict[str, JsonValue] | None,
-    BeforeValidator(_coerce_metadata),
-]
-
-
 class _BaseFeatureSpec(FrozenBaseModel):
     key: Annotated[FeatureKey, BeforeValidator(FeatureKeyAdapter.validate_python)]
     deps: list[FeatureDep] | None = None
@@ -241,8 +220,8 @@ class _BaseFeatureSpec(FrozenBaseModel):
             )
         ]
     )
-    metadata: MetadataField = pydantic.Field(
-        default=None,
+    metadata: dict[str, JsonValue] = pydantic.Field(
+        default_factory=dict,
         description="Metadata attached to this feature.",
     )
 
@@ -258,7 +237,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: list[str] | None = None,
-        metadata: Mapping[str, Any] | JsonValue | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from string key."""
         ...
@@ -271,7 +250,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: list[str] | None = None,
-        metadata: Mapping[str, Any] | JsonValue | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from sequence of parts."""
         ...
@@ -284,7 +263,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: list[str] | None = None,
-        metadata: Mapping[str, Any] | JsonValue | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from FeatureKey instance."""
         ...
@@ -297,11 +276,12 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: list[str] | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from BaseFeatureSpec instance."""
         ...
 
-    def __init__(self, key: CoercibleToFeatureKey | Self, **kwargs):
+    def __init__(self, key: CoercibleToFeatureKey | Self, **kwargs: Any):
         if isinstance(key, type(self)):
             key = key.key
         else:
@@ -384,8 +364,6 @@ def feature_spec_version(self) -> str:
         # Use model_dump with mode="json" for deterministic serialization
         # This ensures all types (like FeatureKey) are properly serialized
         spec_dict = self.model_dump(mode="json")
-        if spec_dict.get("metadata") == {}:
-            spec_dict.pop("metadata", None)
 
         # Sort keys to ensure deterministic ordering
         spec_json = json.dumps(spec_dict, sort_keys=True)
@@ -422,6 +400,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: Sequence[str] | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from string key."""
         ...
@@ -434,6 +413,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: Sequence[str] | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from sequence of parts."""
         ...
@@ -446,6 +426,7 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: Sequence[str] | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from FeatureKey instance."""
         ...
@@ -458,11 +439,12 @@ def __init__(
         deps: list[FeatureDep] | None = None,
         fields: list[FieldSpec] | None = None,
         id_columns: Sequence[str] | None = None,
+        metadata: Mapping[str, JsonValue] | None = None,
     ) -> None:
         """Initialize from FeatureSpec instance."""
         ...
 
-    def __init__(self, key: CoercibleToFeatureKey | Self, **kwargs):
+    def __init__(self, key: CoercibleToFeatureKey | Self, **kwargs: Any):
         # id_columns is always set for FeatureSpec
         super().__init__(key=key, **kwargs)
 

tests/__snapshots__/test_feature_project_detection.ambr

@@ -11,7 +11,7 @@
 # name: test_feature_project_persists_across_graph_operations
   dict({
     'project': 'persist_project',
-    'tracking_version': '2a72c9e2',
+    'tracking_version': '47e9595c',
   })
 # ---
 # name: test_multiple_features_same_project

tests/__snapshots__/test_feature_tracking_version.ambr

@@ -12,8 +12,8 @@
     'feature_versions_same': True,
     'project_a': 'project_a',
     'project_b': 'project_b',
-    'tracking_version_a': '7da02a89202f8eb9f19782512fadf988b815df8255e52a1f35dec9da76e18b12',
-    'tracking_version_b': '5908b5019bdc1acd103482f1d144338f407f50d9d1aa09ca3808fe376f5ffd98',
+    'tracking_version_a': '51105d501a6e001dcbb7c3edb9e4ce93b12a26dc1d0b445cca1de9e986edd3a1',
+    'tracking_version_b': 'e0246e418b8079fd6568fe80959809d0c4161b3c8e4ce70a0e4ecb0adeee0d1b',
     'tracking_versions_differ': True,
   })
 # ---
@@ -43,5 +43,5 @@
   })
 # ---
 # name: test_tracking_version_deterministic
-  '139c7982886b761076245c878e6ad545f4afb9ffabab9238d7af883aff161b9b'
+  '93c7499560ffbab47bb760edd62a6c6be19c7d3ac6977b0286ab5eea4716c23a'
 # ---

tests/__snapshots__/test_id_columns.ambr

@@ -1,8 +1,8 @@
 # serializer version: 1
 # name: test_snapshot_stability_with_id_columns
   dict({
-    'composite': '8dc629cb2bc4bf404e15495f03c7efdcb339b96531fe6579d3f927a82589f843',
-    'default': '0f5dac77217aa051fd504566de0fa3e1d8d4a360ed7958e716ff0f6ebb483a34',
-    'single_custom': 'cfca9196fcb832b9da472042d9fd1ae5c485352d8e409d3d7c1f739880249c22',
+    'composite': 'd51aec13661dd2d2aea4a832b894c93d73ea5d4f5173202ec03ea2e1f483fc6e',
+    'default': '2f4f9c1127afb2868dcd21018adcbfd3680f4eab9e9fb9aefe96d6c9f7a1bd24',
+    'single_custom': 'c571093281edaecb612ae310a509e20047688f00016d1e8fdf0c01361144f244',
   })
 # ---

tests/__snapshots__/test_spec_version.ambr

@@ -1,14 +1,14 @@
 # serializer version: 1
 # name: test_feature_spec_version_deterministic
-  'b0741037ecd180b5948761b9d73b4dd78efa9cf7224240d3e3b718bae26e5852'
+  '90179e7bd99745bd2494f695fa18c15937353a3788c9b5c4650a81a071c1f591'
 # ---
 # name: test_feature_spec_version_includes_all_properties
-  '8b54ae3cd476f3071a9437d4626658c15e40005380a0e2093585d442866e167b'
+  'e7c76b7981649ad1e0d9c8f373043eb42ef75feb7f1a5e222236afb632f12e97'
 # ---
 # name: test_feature_spec_version_recorded_in_metadata_store
   dict({
     'feature_key': 'recorded/feature',
-    'feature_spec_version': '39bcce7c7d9bc23f78982655717dc56879265b87304a96f371c7935df5c59d54',
+    'feature_spec_version': 'd64c6d4dcd782fda8988b6f789367a2c10aed2732e350eb3e2a285051c73d493',
     'feature_version': '0c7b2d83252fbf2f689bec7d37c4f7ffb103755d11eafabc1e000a372f415d83',
     'snapshot_version': '14d4294da40f4ab27ca85eb739fffd70e92c92522dbcba995cfa2aa343988bf3',
   })
@@ -34,17 +34,18 @@
         'snapshot',
         'test',
       ]),
-      'metadata': None,
+      'metadata': dict({
+      }),
     }),
-    'feature_spec_version': 'ceaebb2a96f45752b2ca7daf8326b893d029cc56b366d7ee67bd4a3667008a03',
-    'feature_tracking_version': '7dac08b916cbef01ec6971f5dd00a1c520c31ed6c09d1fd2223e82ed418546fb',
+    'feature_spec_version': '6dc07209be338e7e86bb47dc3f376f3e7447efe35510cac9119ee0878c8c20a2',
+    'feature_tracking_version': 'af52f7692de124d551c0753a169b2a0b4a57003818ed08474c6d9c17480b7e20',
     'feature_version': '7cfde77960e3cf10327e9cb97f311418b101fc5ee2c146922a997306e318edaf',
     'project': 'test',
   })
 # ---
 # name: test_feature_spec_version_with_column_selection_and_rename
-  'a412fd83810799e529173eed40fbcc29ed9adaed8c37ac4ffb21edd4f47c60d1'
+  'cf22f9791d5b83dc084c493e52914a21c6fe27d261e78ec03fe22e033f750b4a'
 # ---
 # name: test_feature_spec_version_with_multiple_complex_deps
-  '09d45d5f79ef381519aad8a30dae3c2bb253f8bc2e61c63114b8ac48fa8efcbc'
+  '20703ec5548c00dbc3f0d634a6791f284d6b1cec4701d014fe5ccaec4deef74b'
 # ---

tests/metadata_stores/__snapshots__/test_snapshot_push.ambr

@@ -29,19 +29,19 @@
     'versions': list([
       dict({
         'feature_key': 'downstream',
-        'feature_spec_version': '20d33a795e3990ca63dbde8a5cd9ee3ab48738c4c1f6375cdcbc638202c505e5',
+        'feature_spec_version': '5596900623b161e09c480056de59559d5c88e8c98718691fe6736988ef716f58',
         'feature_version': '7b8e38a11f800fb93c18c0e04c4609d7b52ec12b62177ba8355531d363f01751',
         'snapshot_version': '1ab4ecd9d7af535b843dc22416dd99d5e430317c0a2bcd01fe83ffc2d7392da8',
       }),
       dict({
         'feature_key': 'upstream',
-        'feature_spec_version': 'a8e3517bbf15ca3ecc14d6300c3c78ace14765963993cfc5ffbee8512c6a7655',
+        'feature_spec_version': '72c4fbce316d0efb1a2f90c2dc45b44a2fa97a09f2397423b51c06b591c289c3',
         'feature_version': '8a2ffeab8da447095c5ee7a77e5635a1e16e7f3605021732f50f7002fa258398',
         'snapshot_version': '1ab4ecd9d7af535b843dc22416dd99d5e430317c0a2bcd01fe83ffc2d7392da8',
       }),
       dict({
         'feature_key': 'downstream',
-        'feature_spec_version': '8ba901bb7efebf6b75ab81334237508b6e399f5515d818e3f76a4cd6421991a4',
+        'feature_spec_version': '2deaa627b51f493d71c80ae0a480c069077a5b1ea9876e55b51bc68f7e514800',
         'feature_version': '7b8e38a11f800fb93c18c0e04c4609d7b52ec12b62177ba8355531d363f01751',
         'snapshot_version': '1ab4ecd9d7af535b843dc22416dd99d5e430317c0a2bcd01fe83ffc2d7392da8',
       }),

tests/test_feature_metadata.py

@@ -1,6 +1,9 @@
 import json
+from typing import Any, cast
 
 import pytest
+from pydantic import ValidationError
+from pydantic.types import JsonValue
 
 from metaxy import Feature, FeatureKey, FeatureSpec
 from metaxy.models.feature import FeatureGraph
@@ -43,31 +46,29 @@ class MetadataFeatureB(
 
 def test_metadata_json_serializable() -> None:
     """Validate JSON serialization enforcement for metadata."""
-    valid_metadata = {
+    valid_metadata: dict[str, JsonValue] = {
         "string": "value",
         "number": 42,
         "float": 3.14,
         "bool": True,
         "null": None,
-        "list": [1, 2, 3],
-        "nested": {"key": "value"},
+        "list": cast(JsonValue, [1, 2, 3]),
+        "nested": cast(JsonValue, {"key": "value"}),
     }
 
     spec = FeatureSpec(
         key=FeatureKey(["tests", "metadata", "json"]),
         deps=None,
         metadata=valid_metadata,
     )
-    assert spec.metadata is not None
-    assert isinstance(spec.metadata, dict)
-    assert isinstance(spec.metadata["list"], list)
+    assert spec.metadata == valid_metadata
     assert json.dumps(spec.metadata) is not None
 
-    with pytest.raises(ValueError):
-        FeatureSpec(
+    with pytest.raises(ValidationError):
+        _ = FeatureSpec(
             key=FeatureKey(["tests", "metadata", "json"]),
             deps=None,
-            metadata={"func": lambda x: x},
+            metadata=cast(Any, {"func": object()}),
         )
 
 
@@ -78,7 +79,16 @@ def test_metadata_immutable() -> None:
         deps=None,
         metadata={"key": "value"},
     )
-    assert spec.metadata is not None
+    assert spec.metadata == {"key": "value"}
 
     with pytest.raises(Exception):
         spec.metadata = {"key": "new_value"}  # type: ignore[assignment]
+
+
+def test_metadata_defaults_to_empty_dict() -> None:
+    """Metadata defaults to an empty dict for easier usage."""
+    spec = FeatureSpec(
+        key=FeatureKey(["tests", "metadata", "default"]),
+        deps=None,
+    )
+    assert spec.metadata == {}