A subclass of dict that filters keys and values

filterdict.py

from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    Hashable,
    Iterable,
    Mapping,
    Optional,
    Tuple,
    TypeVar,
    Union,
    overload,
)

__all__ = ("FilterDict",)

K = TypeVar("K", bound=Optional[Hashable])
V = TypeVar("V")

InitIterable = Iterable[Tuple[K, V]]
# Better name?
InitSequence = Union[Mapping[K, V], InitIterable[K, V]]


class FilterDict(Dict[K, V], Generic[K, V]):
    """
    TODO: docstring :PppPppP
    """

    @overload
    def _filter_seq_and_kwargs(
        self, seq: None, kwargs: Mapping[str, V]
    ) -> Tuple[None, Dict[str, V]]:
        ...

    @overload
    def _filter_seq_and_kwargs(
        self, seq: InitSequence[K, V], kwargs: Mapping[str, V]
    ) -> Tuple[Dict[K, V], Dict[str, V]]:
        ...

    def _filter_seq_and_kwargs(self, seq, kwargs):
        # Filter kwargs
        kwargs = {k: v for k, v in kwargs if not self.filter(k, v)}

        # Filter sequence
        if seq is None:
            mapping = None
        else:
            items = seq.items() if isinstance(seq, Mapping) else seq
            mapping = {k: v for k, v in items if not self.filter(k, v)}

        return mapping, kwargs

    def __init__(
        self,
        filter_: Callable[[K, V], bool],
        /,
        seq: InitSequence[K, V] = (),
        **kwargs: V,
    ) -> None:
        self.filter = filter_
        seq, kwargs = self._filter_seq_and_kwargs(seq, kwargs)
        super().__init__(seq, **kwargs)

    def __setitem__(self, key: K, value: V) -> None:
        if self.filter(key, value):
            try:
                del self[key]
            except KeyError:
                pass
        else:
            super().__setitem__(key, value)

    @overload
    def update(self, __m: Mapping[K, V], **kwargs: V) -> None:
        ...

    @overload
    def update(self, __m: InitIterable[K, V], **kwargs: V) -> None:
        ...

    @overload
    def update(self, **kwargs: V) -> None:
        ...

    def update(self, __m=None, **kwargs) -> None:
        seq, kwargs = self._filter_seq_and_kwargs(__m, kwargs)
        if seq is None:
            super().update(**kwargs)
        else:
            super().update(seq, **kwargs)

    def setdefault(self, k: K, default: V = ...) -> V:
        raise NotImplementedError

test_filterdict.py

from unittest import TestCase

from utils.filterdict import FilterDict


class TestFilterDict(TestCase):
    def test_filters_value_on_instantiation(self):
        d = FilterDict(lambda _, v: v is None, {"a": None})
        self.assertEqual(d, {})

    def test_filters_value_on_set_item(self):
        d = FilterDict(lambda _, v: v is None)
        d["a"], d["b"] = "a", None
        self.assertEqual(d, {"a": "a"})

    def test_filters_value_on_update(self):
        d = FilterDict(lambda _, v: v is None)
        d.update({"a": None})
        self.assertEqual(d, {})

    def test_filters_key_on_instantiation(self):
        d = FilterDict(lambda k, _: k is None, {None: "a"})
        self.assertEqual({}, d)

    def test_filters_key_on_set_item(self):
        d = FilterDict(lambda k, _: k is None)
        d[None] = "a"
        self.assertEqual(d, {})

    def test_filters_key_on_update(self):
        d = FilterDict(lambda k, _: k is None)
        d.update({None: "a"})
        self.assertEqual(d, {})

    def test_set_item_deletes_on_filter_match(self):
        d: dict = FilterDict(lambda _, v: v is None, {"a": "a"})
        d["a"] = None
        self.assertEqual(d, {})

No filter is applied when calling .update():

>>> d = FilterDict(lambda v: v is None)
>>> d.update({"key": None})
>>> d
{"key": None}

Something like this is missing:

def update(self, *args: Any, **kwargs: Any) -> None:
    super().update(*args, **kwargs)
    violating_keys = tuple(k for k, v in self.items() if self.filter(v))
    for k in violating_keys:
        del self[k]

(It would probably be nicer to modify args and/or kwargs before calling super though)

@flaeppe Yeah, I haven't checked that this covers every mutating method.