Simplest way to subclass Set in a way that modifies Set methods?-CodePudding

I want to subclass set in Python so that it behaves as much as a regular set as possible, with a small modification, as shown below.

from typing import Iterable, Any
from collections.abc import MutableSet


class MSet(MutableSet):
    """ Set with a filename attribute """
    def __init__(self, fname: str, iterable: Iterable = None):
        self._fname = fname
        if iterable is None:
            self.elements = set()
        else:
            self.elements = set(iterable)

    def __repr__(self) -> str:
        return f"NAME:{self._fname}; VALUE:{self.elements}"

    def __contains__(self, value: Any) -> bool:
        return value in self.elements

    def __iter__(self) -> Iterable:
        return iter(self.elements)

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

    def add(self, value: Any) -> None:
        self.elements.add(value)

    def discard(self, value: Any) -> None:
        self.elements.discard(value)

    @property
    def filename(self) -> str:
        return self._fname

This is fine, but I want to be able to make it so that the double-under methods

__le__, __lt__, __eq__, __ne__, __gt__, __ge__, __and__, __or__, __sub__, __xor__, 
isdisjoint, __ior__, __iand__, __ixor__, and __isub__

only work when two MSet instances share the same "self._fname". In other words, if I want to make a comparison between mset_a, and mset_b, where their file names are "a.txt" and "b.txt", I want the following:

mset_a = MSet("a.txt", [1, 2, 3])
mset_b = MSet("b.txt", [2, 3, 4])

mset_a & mset_b  # raise exception

How do I go about implementing this?

CodePudding user response：

You need to implement the __and__ method:

def __and__(self, other):
    if self.filename != other.filename:
        raise Exception()
    return self.elements & other.elements