test: FrozenSchema cov

Author: dangotbanned

narwhals/_plan/schema.py

@@ -4,15 +4,15 @@
 from functools import lru_cache
 from itertools import chain
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar, final, overload
 
 from narwhals._plan._expr_ir import NamedIR
 from narwhals._plan._immutable import Immutable
 from narwhals._utils import _hasattr_static
 from narwhals.dtypes import Unknown
 
 if TYPE_CHECKING:
-    from collections.abc import ItemsView, Iterator, KeysView, ValuesView
+    from collections.abc import ItemsView, Iterable, Iterator, KeysView, ValuesView
 
     from typing_extensions import Never, TypeAlias, TypeIs
 
@@ -22,7 +22,7 @@
 
 
 IntoFrozenSchema: TypeAlias = (
-    "IntoSchema | Iterator[tuple[str, DType]] | FrozenSchema | HasSchema"
+    "IntoSchema | Iterable[tuple[str, DType]] | FrozenSchema | HasSchema"
 )
 """A schema to freeze, or an already frozen one.
 
@@ -35,14 +35,15 @@
 _T2 = TypeVar("_T2")
 
 
+@final
 class FrozenSchema(Immutable):
     """Use `freeze_schema(...)` constructor to trigger caching!"""
 
     __slots__ = ("_mapping",)
     _mapping: MappingProxyType[str, DType]
 
-    def __init_subclass__(cls, *_: Never, **__: Never) -> Never:  # pragma: no cover
-        msg = f"Cannot subclass {cls.__name__!r}"
+    def __init_subclass__(cls, *_: Never, **__: Never) -> Never:
+        msg = f"Cannot subclass {FrozenSchema.__name__!r}"
         raise TypeError(msg)
 
     def merge(self, other: FrozenSchema, /) -> FrozenSchema:
@@ -111,7 +112,7 @@ def items(self) -> ItemsView[str, DType]:
     def keys(self) -> KeysView[str]:
         return self._mapping.keys()
 
-    def values(self) -> ValuesView[DType]:  # pragma: no cover
+    def values(self) -> ValuesView[DType]:
         return self._mapping.values()
 
     @overload
@@ -121,15 +122,15 @@ def get(self, key: str, default: DType | _T2, /) -> DType | _T2: ...
     def get(self, key: str, default: DType | _T2 | None = None, /) -> DType | _T2 | None:
         if default is not None:
             return self._mapping.get(key, default)
-        return self._mapping.get(key)  # pragma: no cover
+        return self._mapping.get(key)
 
     def __iter__(self) -> Iterator[str]:
         yield from self._mapping
 
-    def __contains__(self, key: object) -> bool:  # pragma: no cover
+    def __contains__(self, key: object) -> bool:
         return self._mapping.__contains__(key)
 
-    def __getitem__(self, key: str, /) -> DType:  # pragma: no cover
+    def __getitem__(self, key: str, /) -> DType:
         return self._mapping.__getitem__(key)
 
     def __len__(self) -> int:

tests/plan/schema_test.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals as nw
+from narwhals._plan.schema import FrozenSchema, freeze_schema
+from tests.plan.utils import dataframe
+
+
+def test_schema() -> None:
+    mapping = {"a": nw.Int64(), "b": nw.String()}
+    schema = nw.Schema(mapping)
+    frozen_schema = freeze_schema(mapping)
+
+    assert frozen_schema.keys() == schema.keys()
+    assert tuple(frozen_schema.values()) == tuple(schema.values())
+
+    # NOTE: Would type-check if `Schema.__init__` didn't make liskov unhappy
+    assert schema == nw.Schema(frozen_schema)  # type: ignore[arg-type]
+    assert mapping == dict(frozen_schema)
+
+    assert frozen_schema == freeze_schema(mapping)
+    assert frozen_schema == freeze_schema(**mapping)
+    assert frozen_schema == freeze_schema(a=nw.Int64(), b=nw.String())
+    assert frozen_schema == freeze_schema(schema)
+    assert frozen_schema == freeze_schema(frozen_schema)
+    assert frozen_schema == freeze_schema(frozen_schema.items())
+
+    # NOTE: Using `**` unpacking, despite not inheriting from `Mapping` or `dict`
+    assert frozen_schema == freeze_schema(**frozen_schema)
+
+    # NOTE: Using `HasSchema`
+    df = dataframe({"a": [1, 2, 3], "b": ["c", "d", "e"]})
+    assert frozen_schema == freeze_schema(df)
+
+    # NOTE: In case this all looks *too good* to be true
+    assert frozen_schema != freeze_schema(**mapping, c=nw.Float64())
+
+    assert frozen_schema["a"] == schema["a"]
+
+    assert frozen_schema.get("c") is None
+    assert frozen_schema.get("c", nw.Unknown) is nw.Unknown
+    assert frozen_schema.get("c", nw.Unknown()) == nw.Unknown()
+
+    assert "b" in frozen_schema
+    assert "e" not in frozen_schema
+
+    with pytest.raises(TypeError, match="Cannot subclass 'FrozenSchema'"):
+
+        class MutableSchema(FrozenSchema): ...  # type: ignore[misc]