Skip to content
Advertisement

Wrapping a class instance by its own method

Suppose we have a python class like below:

class Test1:
    def __init__(self, value):
        self.value = value

    def f(self):
        return Test(self.value+1)

    def g(self):
        return Test(self.value*2)

test1 = Test1()

Both f and g serves as functions on the class instance in some mathematical notation. For example, I desire to run f(g(value)) in the notation. In the above class definition, the way to achieve it is by running test1.g().f(). Although there is nothing wrong with the code, I just don’t like the fact that the order has changed, and it makes the line not that readable from a mathematics perspective.

I wonder if I could run the line like f(g(value)) with some custom magic or dunder method, but could not find one.

I understand that the way to achieve it is simply taking f and g out from the class like:

class Test1:
    def __init__(self, value):
        self.value = value

def f(obj):
    return Test1(obj.value+1)

def g(obj):
     return Test1(obj.value*2)

test1 = Test1()
new_test = f(g(test1)

However, I have some other similar classes that also have the same function/method names (f, g). So, if I run

from Test1 import Test1, f, g
from Test2 import Test2, f, g

there will be obvious conflicts. I understand that I could do things like

from Test1 import Test1, f as f1, g as g1
from Test2 import Test2, f as f2, g as g2

but, not that pleasant with this.

Is there any custom package like functools that provide any ways? What I am expecting is something like

class Test1:
    def __init__(self, value):
        self.value = value

    @some_special_stuff
    def f(self):
        return Test(self.value+1)

    @some_special_stuff
    def g(self):
        return Test(self.value*2)

test1 = Test1()
new_test = f(g(test1)

Thank you in advance.

Advertisement

Answer

Your Test1 and Test2 classes have the same interface. We can model that

from typing import Protocol

# using a protocol here, so we don't have to explicitly inherit from TestProtocol
class TestProtocol(Protocol):
    value: int
    def f(self) -> 'TestProtocol': ...
    def g(self) -> 'TestProtocol': ...

Now we know the interface and can define the free function f and g in terms of that interface:

def f(test: TestProtocol) -> TestProtocol:
    return test.f()


def g(test: TestProtocol) -> TestProtocol:
    return test.g()

Now the free functions f and g are compatible with every class that implements the protocol:

class Test1:
    def __init__(self, value: int) -> None:
        self.value = value
    
    def f(self) -> 'Test1':
        return Test1(self.value + 1)
    
    def g(self) -> 'Test1':
        return Test1(self.value * 2)


class Test2:
    def __init__(self, value: int) -> None:
        self.value = value
    
    def f(self) -> 'Test2':
        return Test2(self.value - 1)
    
    def g(self) -> 'Test2':
        return Test2(self.value // 2)


print(f(g(Test1(4))).value)  # 9
print(f(g(Test2(4))).value)  # 1

Actually, you don’t need the protocol at all. If you don’t care about mypy, remove all type hints and the protocol.

So here is the “short” form without type hints:

class Test1:
    def __init__(self, value):
        self.value = value
    
    def f(self):
        return Test1(self.value + 1)
    
    def g(self):
        return Test1(self.value * 2)

class Test2:
    ...

def f(test):
    return test.f()


def g(test):
    return test.g()


print(f(g(Test1(4))).value)
print(f(g(Test2(4))).value)
User contributions licensed under: CC BY-SA
6 People found this is helpful
Advertisement