feat: add some dunder methods to Schema

src/safeds/data/tabular/typing/_schema.py

@@ -1,22 +1,20 @@
 from __future__ import annotations
 
 import sys
-from typing import TYPE_CHECKING
+from collections.abc import Iterator, Mapping
+from typing import TYPE_CHECKING, Any
 
 from safeds._utils import _structural_hash
 from safeds._validation import _check_columns_exist
 
+from ._column_type import ColumnType
 from ._polars_column_type import _PolarsColumnType
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
-
     import polars as pl
 
-    from ._column_type import ColumnType
-
 
-class Schema:
+class Schema(Mapping[str, ColumnType]):
     """The schema of a row or table."""
 
     # ------------------------------------------------------------------------------------------------------------------
@@ -41,16 +39,28 @@ def __init__(self, schema: Mapping[str, ColumnType]) -> None:
             check_dtypes=False,
         )
 
+    def __contains__(self, name: Any) -> bool:
+        return self.has_column(name)
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Schema):
             return NotImplemented
         if self is other:
             return True
         return self._schema == other._schema
 
+    def __getitem__(self, key: str, /) -> ColumnType:
+        return self.get_column_type(key)
+
     def __hash__(self) -> int:
         return _structural_hash(tuple(self._schema.keys()), [str(type_) for type_ in self._schema.values()])
 
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._schema.keys())
+
+    def __len__(self) -> int:
+        return self.column_count
+
     def __repr__(self) -> str:
         return f"Schema({self!s})"
 
@@ -108,7 +118,7 @@ def column_names(self) -> list[str]:
 
     def get_column_type(self, name: str) -> ColumnType:
         """
-        Get the type of a column.
+        Get the type of a column. This is equivalent to using the `[]` operator (indexed access).
 
         Parameters
         ----------
@@ -131,14 +141,17 @@ def get_column_type(self, name: str) -> ColumnType:
         >>> schema = Schema({"a": ColumnType.int64(), "b": ColumnType.float32()})
         >>> schema.get_column_type("a")
         int64
+
+        >>> schema["b"]
+        float32
         """
         _check_columns_exist(self, name)
 
         return _PolarsColumnType(self._schema[name])
 
     def has_column(self, name: str) -> bool:
         """
-        Check if the table has a column with a specific name.
+        Check if the schema has a column with a specific name.  This is equivalent to using the `in` operator.
 
         Parameters
         ----------
@@ -148,7 +161,7 @@ def has_column(self, name: str) -> bool:
         Returns
         -------
         has_column:
-            Whether the table has a column with the specified name.
+            Whether the schema has a column with the specified name.
 
         Examples
         --------
@@ -157,7 +170,7 @@ def has_column(self, name: str) -> bool:
         >>> schema.has_column("a")
         True
 
-        >>> schema.has_column("c")
+        >>> "c" in schema
         False
         """
         return name in self._schema

tests/safeds/data/tabular/typing/_schema/test_contains.py

@@ -0,0 +1,16 @@
+import pytest
+
+from safeds.data.tabular.typing import ColumnType, Schema
+
+
+@pytest.mark.parametrize(
+    ("schema", "column", "expected"),
+    [
+        (Schema({}), "C", False),
+        (Schema({"A": ColumnType.null()}), "A", True),
+        (Schema({"A": ColumnType.null()}), "B", False),
+    ],
+    ids=["empty", "has column", "doesn't have column"],
+)
+def test_should_check_if_column_is_in_schema(schema: Schema, column: str, expected: bool) -> None:
+    assert (column in schema) == expected

tests/safeds/data/tabular/typing/_schema/test_get_item.py

@@ -0,0 +1,29 @@
+import pytest
+
+from safeds.data.tabular.typing import ColumnType, Schema
+from safeds.exceptions import ColumnNotFoundError
+
+
+@pytest.mark.parametrize(
+    ("schema", "name", "expected"),
+    [
+        (
+            Schema({"col1": ColumnType.int64()}),
+            "col1",
+            ColumnType.int64(),
+        ),
+        (
+            Schema({"col1": ColumnType.string()}),
+            "col1",
+            ColumnType.string(),
+        ),
+    ],
+    ids=["int column", "string column"],
+)
+def test_should_return_data_type_of_column(schema: Schema, name: str, expected: ColumnType) -> None:
+    assert schema[name] == expected
+
+
+def test_should_raise_if_column_name_is_unknown() -> None:
+    with pytest.raises(ColumnNotFoundError):
+        _ignored = Schema({})["col1"]

tests/safeds/data/tabular/typing/_schema/test_has_column.py

@@ -12,5 +12,5 @@
     ],
     ids=["empty", "has column", "doesn't have column"],
 )
-def test_should_return_if_column_is_in_schema(schema: Schema, column: str, expected: bool) -> None:
+def test_should_check_if_column_is_in_schema(schema: Schema, column: str, expected: bool) -> None:
     assert schema.has_column(column) == expected

tests/safeds/data/tabular/typing/_schema/test_iter.py

@@ -0,0 +1,20 @@
+import pytest
+
+from safeds.data.tabular.typing import ColumnType, Schema
+
+
+@pytest.mark.parametrize(
+    ("schema", "expected"),
+    [
+        (Schema({}), []),
+        (Schema({"col1": ColumnType.null()}), ["col1"]),
+        (Schema({"col1": ColumnType.null(), "col2": ColumnType.null()}), ["col1", "col2"]),
+    ],
+    ids=[
+        "empty",
+        "one column",
+        "two columns",
+    ],
+)
+def test_should_return_column_names(schema: Schema, expected: list[str]) -> None:
+    assert list(schema) == expected

tests/safeds/data/tabular/typing/_schema/test_len.py

@@ -0,0 +1,20 @@
+import pytest
+
+from safeds.data.tabular.typing import ColumnType, Schema
+
+
+@pytest.mark.parametrize(
+    ("schema", "expected"),
+    [
+        (Schema({}), 0),
+        (Schema({"col1": ColumnType.null()}), 1),
+        (Schema({"col1": ColumnType.null(), "col2": ColumnType.null()}), 2),
+    ],
+    ids=[
+        "empty",
+        "one column",
+        "two columns",
+    ],
+)
+def test_should_return_number_of_columns(schema: Schema, expected: int) -> None:
+    assert len(schema) == expected