from collections import OrderedDict
from typing import Callable, TypeVar, Union, Generic, Optional, Iterable, ItemsView, KeysView, ValuesView
T = TypeVar("T")
[docs]class ElementSet(Generic[T]):
"""
Strictly typed collection that mostly resembles an ordered `set`, but borrows some behavior of `list` and `dict`.
Items must be representable by a `str` key (like `item.name`), allowing `dict`-like access.
"""
def __init__(self, key_func: Callable[[T], str], check_func: Callable[[T], bool],
items: Optional[Iterable[T]] = None):
"""
Create an empty ElementSet.
:param key_func: Function to get a string key from an item (e.g. `lambda item: item.name`)
:param check_func: Function to check if an item is valid for this ElementSet
(e.g. `lambda item: hasattr(item, 'name')` or `lambda item: isinstance(item, ...)`).
If it returns `True`, everything is fine. If it returns `False`, a default exception is
raised. To raise a more specific exception, raise an exception instead of returning `False`.
"""
self.__dict: "OrderedDict[str, T]" = OrderedDict()
self.__key_func = key_func
self.__check_func = check_func
if items is not None:
self.update(items)
[docs] def add(self, item: T, as_first: bool = False) -> None:
"""Add an item to the collection. If `as_first`, the item is added at the beginning, else at the end."""
if self.__check_func(item):
key = self.__key_func(item)
self.__dict[key] = item
if as_first:
self.__dict.move_to_end(key, last=False)
else:
raise TypeError(f"Invalid item: {item}")
[docs] def remove(self, item: Union[str, T]) -> None:
"""Remove the item. If key is a string, remove the item specified by `key`"""
if isinstance(item, str):
self.__dict.pop(item)
elif self.__check_func(item):
if item in self:
self.__dict.pop(self.__key_func(item))
else:
raise ValueError(f"Item {item!r} not found")
else:
raise TypeError(f"Expected a string key or item of correct type, got {item}")
[docs] def keys(self) -> KeysView[str]:
"""Like `dict.keys`"""
return self.__dict.keys()
[docs] def items(self) -> ItemsView[str, T]:
"""Like `dict.items`"""
return self.__dict.items()
def __getitem__(self, key: Union[str, int]) -> T:
"""Get the item specified by `key`. The `key` can be `str` key like in `dict` or `int` index like in `list`."""
if isinstance(key, str):
return self.__dict[key]
if isinstance(key, int):
return tuple(self.__dict.values())[key]
raise KeyError(f"Expected an integer index or string key, got {type(key)}")
def __len__(self) -> int:
return len(self.__dict)
[docs] def clear(self) -> None:
"""Like `dict.clear`"""
self.__dict.clear()
[docs] def copy(self) -> "ElementSet[T]":
"""Return a shallow copy of this collection"""
cpy = self.__class__(key_func=self.__key_func, check_func=self.__check_func)
cpy.update(self)
return cpy
[docs] def values(self) -> ValuesView[T]:
"""Like `dict.values`"""
return self.__dict.values()
def __iter__(self) -> Iterable[T]:
"""Iterate over the stored items"""
return iter(self.values())
[docs] def update(self, *items: Union[T, Iterable[T]]) -> None:
"""Like `set.update`"""
if len(items) == 1 and not self.__check_func(items[0]):
# e.g. self.update([1, 2, 3])
for item in items[0]:
self.add(item)
else:
# e.g. self.update(1, 2, 3)
for item in items:
self.add(item)
[docs] def pop(self, key: Optional[Union[int, str]] = None) -> T:
"""Like `dict.pop` if `key` is `str`, else like `list.pop`"""
if key is None:
obj = self[-1]
else:
obj = self[key]
self.remove(obj)
return obj
def __contains__(self, item: Union[T, str]) -> bool:
"""For string keys: True if key is in self.keys(), else True if key in self.values()"""
if isinstance(item, str):
return item in self.__dict.keys()
elif self.__check_func(item):
return item in self.__dict.values()
raise TypeError(f"Expected a string key or item of correct type, got {type(item)}")
def __repr__(self) -> str:
return f"{self.__class__.__name__}({[item for item in self]!r})"
def __eq__(self, other):
if isinstance(other, ElementSet):
return dict(self.__dict) == dict(other.__dict) # convert to dict for unordered comparison
return False