fix: Use covariant type annotation for collection sampling type

Author: borchero

dataframely/collection.py

@@ -4,7 +4,7 @@
 import sys
 import warnings
 from abc import ABC
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any, Generic, Self, TypeVar, cast
 
@@ -20,10 +20,10 @@
 
 if sys.version_info >= (3, 13):
     SamplingType = TypeVar(
-        "SamplingType", bound=MutableMapping[str, Any], default=dict[str, Any]
+        "SamplingType", bound=Mapping[str, Any], default=dict[str, Any]
     )
 else:  # pragma: no cover
-    SamplingType = TypeVar("SamplingType", bound=MutableMapping[str, Any])
+    SamplingType = TypeVar("SamplingType", bound=Mapping[str, Any])
 
 
 class Collection(BaseCollection, ABC, Generic[SamplingType]):

tests/test_typing.py

@@ -9,7 +9,7 @@
 import datetime
 import decimal
 import functools
-from typing import Any
+from typing import Any, NotRequired, TypedDict
 
 import polars as pl
 import pytest
@@ -61,7 +61,21 @@ class MySecondSchema(dy.Schema):
     b = dy.Integer()
 
 
-class MyCollection(dy.Collection):
+class SamplingTypeFirst(TypedDict):
+    a: NotRequired[int]
+
+
+class SamplingTypeSecond(TypedDict):
+    a: NotRequired[int]
+    b: NotRequired[int]
+
+
+class SamplingType(TypedDict):
+    first: NotRequired[SamplingTypeFirst]
+    second: NotRequired[SamplingTypeSecond]
+
+
+class MyCollection(dy.Collection[SamplingType]):
     first: dy.LazyFrame[MyFirstSchema]
     second: dy.LazyFrame[MySecondSchema]