feat(datatype): implement Mapping abstract base class for StructType

Author: kszucs

ibis/expr/datatypes/core.py

@@ -2,6 +2,7 @@
 
 import numbers
 from abc import abstractmethod
+from collections.abc import Iterator, Mapping
 from typing import Any, Iterable, NamedTuple
 
 import numpy as np
@@ -639,7 +640,7 @@ def to_integer_type(self):
 
 
 @public
-class Struct(DataType):
+class Struct(DataType, Mapping):
     """Structured values."""
 
     fields = frozendict_of(instance_of(str), datatype)
@@ -677,6 +678,12 @@ def types(self) -> tuple[DataType, ...]:
         """Return the types of the struct's fields."""
         return tuple(self.fields.values())
 
+    def __len__(self) -> int:
+        return len(self.fields)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.fields)
+
     def __getitem__(self, key: str) -> DataType:
         return self.fields[key]
 

ibis/tests/expr/test_datatypes.py

@@ -141,6 +141,45 @@ def test_struct_with_string_types():
     )
 
 
+def test_struct_mapping_api():
+    s = dt.Struct(
+        {
+            'a': 'map<double, string>',
+            'b': 'array<map<string, array<int32>>>',
+            'c': 'array<string>',
+            'd': 'int8',
+        }
+    )
+
+    assert s['a'] == dt.Map(dt.double, dt.string)
+    assert s['b'] == dt.Array(dt.Map(dt.string, dt.Array(dt.int32)))
+    assert s['c'] == dt.Array(dt.string)
+    assert s['d'] == dt.int8
+
+    assert 'a' in s
+    assert 'e' not in s
+    assert len(s) == 4
+    assert tuple(s) == s.names
+    assert tuple(s.keys()) == s.names
+    assert tuple(s.values()) == s.types
+    assert tuple(s.items()) == tuple(zip(s.names, s.types))
+
+    s1 = s.copy()
+    s2 = dt.Struct(
+        {
+            'a': 'map<double, string>',
+            'b': 'array<map<string, array<int32>>>',
+            'c': 'array<string>',
+        }
+    )
+    assert s == s1
+    assert s != s2
+
+    # doesn't support item assignment
+    with pytest.raises(TypeError):
+        s['e'] = dt.int8
+
+
 @pytest.mark.parametrize(
     'case',
     [