Decorator to do a test before launching each function

Question

Instead of repeating if not self.connected: ... return for every method of this class:

class A:
    def __init__(self):
        self.connected = False
    def connect(self, password):
        if password == "ab":
            self.connected = True
    def send(self, message):
        if not self.connected:
            print("abandoned because not connected")
            return
        print("sending")
    def receive(self):
        if not self.connected:
            print("abandoned because not connected")
            return
        print("receiving")
a = A()
a.send()

is it possible to do it with a decorator like

class A:

    ...

    def requires_connected(self, func):
        if self.connected:
            func()
        else:
            print("abandoned because not connected")

    @requires_connected
    def send(self, message):
        print("sending")

?

The latter doesn't work (TypeError: requires_connected() missing 1 required positional argument: 'func'), probably because of the arguments.

How to make this decorator work correctly?

Answer 1

Using Declaring decorator inside a class, I finally found the solution:

class A:
    def __init__(self):
        self.connected = False

    def requires_connected(func):
        def wrapper(self, *args, **kwargs):
            if self.connected:
                return func(self, *args, **kwargs)
            else:
                print("abandoned because not connected")
        return wrapper

    @requires_connected
    def send(self, message):
        print(f"sending '{message}'")

    # this is syntactic sugar for:
    #    def send(self, message):
    #       ...
    #    send = requires_connected(send)

a = A()
a.send("hello")
a.connected = True
a.send("hello")