Skip to content

Commit

Permalink
feat(datatype): implement Mapping abstract base class for StructType
Browse files Browse the repository at this point in the history
  • Loading branch information
kszucs committed Jan 30, 2023
1 parent a8d3007 commit 5df2022
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 1 deletion.
9 changes: 8 additions & 1 deletion ibis/expr/datatypes/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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]

Expand Down
39 changes: 39 additions & 0 deletions ibis/tests/expr/test_datatypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
[
Expand Down

0 comments on commit 5df2022

Please sign in to comment.