operators.py

import operator
import inspect
import pytest


def getattribute(obj, name):
    return object.__getattribute__(obj, name)

def setattribute(obj, name, value):
    object.__setattr__(obj, name, value)


def get_n_args(f) -> int:
    return len(inspect.getfullargspec(f).args)


OPERATOR_SLOT_METHODS = {
    n: get_n_args(f)
    for n, f in vars(operator).items()
    if n.startswith("__") and getattr(f, "__name__", None) in operator.__all__
}
OPERATOR_SLOT_METHODS.update(
    {f"__r{n[3:]}": j for n, j in OPERATOR_SLOT_METHODS.items() if n.startswith("__i")}
)


class Proxy:
    def __init__(self, value):
        setattribute(self, 'value', value)

    def __setattr__(self, key, value):
        wrapped = getattribute(self, 'value')
        setattr(wrapped, key, value)

    def __getattribute__(self, item):
        wrapped = getattribute(self, "value")
        return type(self)(getattr(wrapped, item))

    for method, n_args in OPERATOR_SLOT_METHODS.items():
        is_in_place = method.startswith("__i")
        method_args = ", ".join([f"arg{i}" for i in range(n_args - 1)])

        if is_in_place:
            method_body = f"""
def {method}(self, {method_args}):
    wrapped = getattribute(self, 'value')
    result = wrapped.{method}({method_args})
    setattribute(self, 'value', result)
    return self
        """

        else:
            method_body = f"""
def {method}(self, {method_args}):
    wrapped = getattribute(self, 'value')
    result = wrapped.{method}({method_args})
    return type(self)(result)
        """
        # print(method_body)
        exec(method_body)

    def __repr__(self):
        cls_name = type(self).__qualname__
        wrapped = getattribute(self, "value")
        return f"{cls_name}({wrapped!r})"


# Syntax 1
class EllipsisSignal(Proxy):
    def __setitem__(self, key, value):
        if key is Ellipsis:
            setattribute(self, 'value', value)
        else:
            super().__setitem__(key, value)


# Syntax 2
class SignalAssignment:
    def __init__(self, value):
        self.value = value


push = SignalAssignment


class ILShiftSignal(Proxy):
    def __ilshift__(self, other):
        if isinstance(other, SignalAssignment):
            setattribute(self, 'value', other.value)
            return self
        else:
            return super().__ilshift__(other)


# Tests
def test_ellipsis():
    e = EllipsisSignal({})

    # Check non-assignment behaviour
    e[1] = 'Jack'
    assert getattribute(e, 'value') == {1: 'Jack'}

    # Check assignment behaviour
    e[...] = {2: 'Jill'}
    assert getattribute(e, 'value') == {2: 'Jill'}


def test_ilshift():
    i = ILShiftSignal(14)

    # Check assignment behaviour
    i <<= push(12)
    assert getattribute(i, 'value') == 12

    # Check non-assignment behaviour
    with pytest.raises(AttributeError):
        i <<= 12