added casted dict class for easy automatic casting

Author: wolph

python_utils/containers.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import abc
+from typing import Any
+from typing import Generator
+
+from . import types
+
+KT = types.TypeVar('KT')
+VT = types.TypeVar('VT')
+DT = types.Dict[KT, VT]
+KT_cast = types.Optional[types.Callable[[Any], KT]]
+VT_cast = types.Optional[types.Callable[[Any], VT]]
+
+# Using types.Union instead of | since Python 3.7 doesn't fully support it
+DictUpdateArgs = types.Union[
+    types.Mapping,
+    types.Iterable[types.Union[types.Tuple[Any, Any], types.Mapping]],
+]
+
+
+class CastedDictBase(types.Dict[KT, VT], abc.ABC):
+    _key_cast: KT_cast
+    _value_cast: VT_cast
+
+    def __init__(
+            self,
+            key_cast: KT_cast = None,
+            value_cast: VT_cast = None,
+            *args,
+            **kwargs
+    ) -> None:
+        self._value_cast = value_cast
+        self._key_cast = key_cast
+        self.update(*args, **kwargs)
+
+    def update(
+            self,
+            *args: DictUpdateArgs,
+            **kwargs
+    ) -> None:
+        if args:
+            kwargs.update(*args)
+
+        if kwargs:
+            for key, value in kwargs.items():
+                self[key] = value
+
+    def __setitem__(self, key: Any, value: Any) -> None:
+        if self._key_cast is not None:
+            key = self._key_cast(key)
+
+        return super().__setitem__(key, value)
+
+
+class CastedDict(CastedDictBase):
+    '''
+    Custom dictionary that casts keys and values to the specified typing.
+
+    Note that you can specify the types for mypy and type hinting with:
+    CastedDict[int, int](int, int)
+
+    >>> d = CastedDict(int, int)
+    >>> d[1] = 2
+    >>> d['3'] = '4'
+    >>> d.update({'5': '6'})
+    >>> d.update([('7', '8')])
+    >>> d
+    {1: 2, 3: 4, 5: 6, 7: 8}
+    >>> list(d.keys())
+    [1, 3, 5, 7]
+    >>> list(d)
+    [1, 3, 5, 7]
+    >>> list(d.values())
+    [2, 4, 6, 8]
+    >>> list(d.items())
+    [(1, 2), (3, 4), (5, 6), (7, 8)]
+    >>> d[3]
+    4
+
+    # Casts are optional and can be disabled by passing None as the cast
+    >>> d = CastedDict()
+    >>> d[1] = 2
+    >>> d['3'] = '4'
+    >>> d.update({'5': '6'})
+    >>> d.update([('7', '8')])
+    >>> d
+    {1: 2, '3': '4', '5': '6', '7': '8'}
+    '''
+
+    def __setitem__(self, key, value):
+        if self._value_cast is not None:
+            value = self._value_cast(value)
+
+        super().__setitem__(key, value)
+
+
+class LazyCastedDict(CastedDictBase):
+    '''
+    Custom dictionary that casts keys and lazily casts values to the specified
+    typing. Note that the values are cast only when they are accessed and
+    are not cached between executions.
+
+    Note that you can specify the types for mypy and type hinting with:
+    LazyCastedDict[int, int](int, int)
+
+    >>> d = LazyCastedDict(int, int)
+    >>> d[1] = 2
+    >>> d['3'] = '4'
+    >>> d.update({'5': '6'})
+    >>> d.update([('7', '8')])
+    >>> d
+    {1: 2, 3: '4', 5: '6', 7: '8'}
+    >>> list(d.keys())
+    [1, 3, 5, 7]
+    >>> list(d)
+    [1, 3, 5, 7]
+    >>> list(d.values())
+    [2, 4, 6, 8]
+    >>> list(d.items())
+    [(1, 2), (3, 4), (5, 6), (7, 8)]
+    >>> d[3]
+    4
+
+    # Casts are optional and can be disabled by passing None as the cast
+    >>> d = LazyCastedDict()
+    >>> d[1] = 2
+    >>> d['3'] = '4'
+    >>> d.update({'5': '6'})
+    >>> d.update([('7', '8')])
+    >>> d
+    {1: 2, '3': '4', '5': '6', '7': '8'}
+    >>> list(d.keys())
+    [1, '3', '5', '7']
+    >>> list(d.values())
+    [2, '4', '6', '8']
+
+    >>> list(d.items())
+    [(1, 2), ('3', '4'), ('5', '6'), ('7', '8')]
+    >>> d['3']
+    '4'
+    '''
+
+    def __setitem__(self, key, value):
+        if self._key_cast is not None:
+            key = self._key_cast(key)
+
+        super().__setitem__(key, value)
+
+    def __getitem__(self, key) -> VT:
+        if self._key_cast is not None:
+            key = self._key_cast(key)
+
+        value = super().__getitem__(key)
+
+        if self._value_cast is not None:
+            value = self._value_cast(value)
+
+        return value
+
+    def items(self) -> Generator[tuple[KT, VT], None, None]:  # type: ignore
+        if self._value_cast is None:
+            yield from super().items()
+        else:
+            for key, value in super().items():
+                yield key, self._value_cast(value)
+
+    def values(self) -> Generator[VT, None, None]:  # type: ignore
+        if self._value_cast is None:
+            yield from super().values()
+        else:
+            for value in super().values():
+                yield self._value_cast(value)
+
+
+if __name__ == '__main__':
+    import doctest
+
+    doctest.testmod()