Recipe for MutableSet

mutable_set.py

"""
filename: mutable_set.py
author: Niraj Kamdar
"""

import weakref
from collections import abc
from typing import Iterator, List, Any


class MutableSet(abc.MutableSet, abc.Mapping):
    def __init__(self, iterable=None):
        self.set = {}
        if iterable:
            self.set = {i: weakref.ref(i) for i in iterable}

    def add(self, x: Any) -> None:
        if x not in self.set:
            self.set[x] = weakref.ref(x)

    def discard(self, x: Any) -> None:
        if x in self.set:
            del self.set[x]

    def remove(self, x: Any) -> None:
        if x not in self.set:
            raise KeyError(x)
        del self.set[x]

    def __contains__(self, x: Any) -> bool:
        return x in self.set

    def __len__(self) -> int:
        return len(self.set)

    def __repr__(self) -> str:
        return f"{{{', '.join(map(str, self.set.keys())).rstrip(', ')}}}"

    def __iter__(self) -> Iterator[Any]:
        return iter(self.set)

    def __getitem__(self, k: Any) -> Any:
        return self.set[k]()
    
    def get(self, k: Any, default=None) -> Any:
        if k in self.set:
            return self.set[k]()
        return default