[Python] Decorate instance methods with decorator class that replaces method with class...

I would like to be able to decorate instance methods with a class Step such that the methods are replaced by a Step object. At the same time, I'd like to have the option be able to instantiate a step with Step(func, options).

Working example


The example below seems to be doing just that but I'm not sure if this a cleaner way of achieving this functionality?

from typing import Callable
import functools


def step(options: "StepOptions" = None):
"""Step decorator"""
def decorator(func: Callable) -> Step:
return Step(func, options)
return decorator


class Step:
def __init__(self, func: Callable, options: "StepOptions" = None):
self.func = func
self.options = options if options else StepOptions()

def __get__(self, instance, owner=None):
# Creating a bound method by passing the instance to functools.partial
bound_method = functools.partial(self.func, instance)
functools.update_wrapper(bound_method, self.func)

# Return a new Step instance with the bound method
return Step(bound_method, self.options)

def __call__(self, *args, **kwargs):
return self.func(*args, **kwargs)

def __repr__(self):
return f"Step(func: {self.func.__name__}, options: {self.options})"

class StepOptions:

def __init__(self, retries: int = 0, timeout_s : float = None):
self.retries = retries
self.timeout_s = timeout_s

def __repr__(self):
return f"StepOptions(retries: {self.retries}, timeout_s: {self.timeout_s})"


class Test:
def __init__(self):
self.steps= [
Step(self.test_1, options=StepOptions(retries=3)),
self.test_2,
self.test_3,
]

self.value = 2

def test_1(self):
print("This is test_1")
return self.value

@step()
def test_2(self) -> str:
print("This is test_2")
return 20

@step(StepOptions(timeout_s=10))
def test_3(self) -> int:
print("This is test_3")
return 30


test = Test()

for test_step in test.steps:
print(f"type: {type(test_step)}")
print(test_step)
result = test_step()
print(result)


which outputs

type: <class '__main__.Step'>
Step(func: test_1, options: StepOptions(retries: 3, timeout_s: None))
This is test_1
2
type: <class '__main__.Step'>
Step(func: test_2, options: StepOptions(retries: 0, timeout_s: None))
This is test_2
20
type: <class '__main__.Step'>
Step(func: test_3, options: StepOptions(retries: 0, timeout_s: 10))
This is test_3
30

Alternative


In particular I'm not sure about the __get__ method of Step. So far I've only seen examples along the lines of

def __get__(self, instance, owner=None):
return functools.partial(self.__call__, instance)


However, when doing that, print(test_2) results in

functools.partial(<bound method Step.__call__ of Step(func: test_2, options: StepOptions(retries: 0, timeout_s: None))>, <__main__.Test object at 0x78f0e1a19fc0>)


instead of

Step(func: test_2, options: StepOptions(retries: 0, timeout_s: None))

Background


A bit of background: I'm planning on using this construct in a test executor. The idea is that by creating a Step instance from each test step function, I have a place to store metadata (such as StepOptions) for each step and access that metadata in the test executor. If anyone sees a better way of doing this I'm more than happy to receive inputs!

Continue reading...