How to customize operator overloading to return custom class instances

  Kiến thức lập trình

I was working on a problem involving numbers in n-base systems. I started working on creating a number class that behaves like an int but is represented as an n-base number.

Here’s how I proceeded:

  1. I built a metaclass to collect method names that are exclusively part of the int type.
  2. I used these names to wrap the int special methods.

However, when I use it as a normal number (e.g., bn + 1), it returns an int. I would like it to return a BaseNumber instance, just like in bn.__add__(1).

Is there a way to make bn + 1 return a BaseNumber instance?

I could simply reimplement each int method manually by wrapping the super() method inside each method:


class _MetaNumberBase(type):

    def __new__(cls, name, bases, dct):
        # Collect method names that are exclusively part of the int type

        cls_attrs = dir(cls) # dir(NumberBase)
        _int_attrs = []
        for base in bases:
            for attr in dir(int):
                if attr not in cls_attrs and attr != "_int_attrs": 
                    # _int_attrs is filtered out since it is declared as a list[str] type in the NumberBase class
                    base_attr = getattr(base, attr)
                    if not inspect.isdatadescriptor(base_attr) and attr.startswith("__"):
                        _int_attrs.append(attr)
                    dct[attr] = base_attr
        # Manually adds the values to NumberBase to communicate them
        dct["_int_attrs"] = _int_attrs
        return super().__new__(cls, name, bases, dct)

class NumberBase(int, metaclass=_MetaNumberBase):
    _int_attrs: list[str]

    def __new__(cls, x: int, base_map, null= None):
        # Create an int instance using __new__
        return super().__new__(cls, x)
    
    def __init__(self, x: int, base_map, null=None) -> None:

        self.base = len(base_map)
        self.base_map = tuple(base_map)
        self._field = dict(enumerate(base_map))

        if null:
            self.null = null
        else:
            self.null = self._field[0]
        
        for attr in dir(self):
            if attr in self._int_attrs:
                method = getattr(self, attr)
                if callable(method):
                    
                    setattr(self, attr, self._decorate(method))
        
        self._immutable = ("base", "base_map", "null", "_immutable", "_locked")
        # This lock is immutable
        self._locked = True

    def to_base(self, base_map, null=None):
        return NumberBase(self, base_map, null)

    def __setattr__(self, name, value):
        if hasattr(self, '_locked') and name in self._immutable:
            raise AttributeError(f"The '{name}' attribute is immutable")
        super().__setattr__(name, value)

    def __delattr__(self, name):
        if name in self._immutable:
            raise AttributeError(f"The '{name}' attribute is immutable")
        super().__delattr__(name)

    def __iter__(self):
        # Converts base 10 to base n
        result = []
        num = self
        if num == 0:
            yield self.null
            return 
        elif self.base == 1:
            result = [self.base_map[0] for x in range(self)]
        else:
            while num > 0:
                remainder = num % self.base
                result.append(self._field[remainder])
                num //= self.base
            result.reverse()
        while result:
            yield result.pop(0)

    def __str__(self) -> str:
        return str().join([str(x) for x in self])

    def __repr__(self):
        return "[bn]: "+str(self)
    
    def _decorate(self, method):
        def wrapper(*args, **kwargs):
            result = method(*args, **kwargs)
            if isinstance(result, NumberBase) or isinstance(result, int):
                return NumberBase(result, base_map=self.base_map, null=self.null)
            else:
                return result
        return wrapper

if __name__ == "__main__":
    bn = NumberBase(34, [0, 1]) # base 2

    print(repr(bn)) # [bn]: 100010
    print(int(bn))  # 34

    # add a number to bn

    s = bn + 2
    # ================== ISSUE ====================
    print(repr(s)) # 36                           | should be [bn]: 100100
    # =============================================
    print(int(s))  # 36

    # add manually a number to bn

    s = bn.__add__(2)
    print(repr(s)) # [bn]: 100100
    print(int(s))  # 36
    

I added 2 to bn like this:

s = bn + 2
assert isinstance(s, bn)
AssertionError

It should do just like:

s = bn.__add__(2)
assert isinstance(s, bn)

3

Theme wordpress giá rẻ Theme wordpress giá rẻ Thiết kế website Kho Theme wordpress Kho Theme WP Theme WP

LEAVE A COMMENT