Search code examples
pythondecoratorclass-decorator

Get instance of class as decorator when the class/decorator has arguments and more methods


In Python you can add metadata to a function as such:

def test(a,b):
    return a+b

test.tags = ["test"]
test.version = "0.1"

print(test.tags) # ['test']

The thing is, I use a decorator to log my functions with an orchestration platform, which supports metadata:

@orchestrator(
    tags = ["test"],
    version = "0.1"
)
def test1(a,b):
    return a+b

Although there are some parts of the metadata that aren't passed to the decorator.

I wanted to create a class that allows me to decorate a function which applies the orchestrator decorator and the metadata as a method (and I may want to add some methods afterwards).

Initially I wrote something like:

class step:

def __init__(self, function, version, description, tags):

    self._function    = function
    self._version     = version
    self._description = description
    self._tags        = tags

    @orchestrator(
        tags = self._tags,
        version = self._version
    )
    def __call__(self):
        return self._function

    def metadata(self):
        return {
            "version": self._version,
            "description": self._description,
            "tags": self._tags,
        }

    # Other methods...

@step(
    tags = ["test"],
    description = "Example description"
    version = "0.1"
)
def test(a,b):
    return a+b

Of course this approach has a couple of problems:

  • That's not how a class as decorator is defined if you need to pass arguments to the class:

     # This:
     @step(
         tags = ["test"],
         description = "Example description"
         version = "0.1"
     )
     def test(a,b):
         return a+b
     # Is the same as this:
     step(tags = ["test"], description = "Example description", version = "0.1")(a,b)
     # So step is missing one arguments: function
    
  • Even if that worked, self._tags, self._version wouldn't be available to the orchestrator decorator if they are outside a method.

So, my next try was:

class step:

    def __init__(self, version, description, tags):

        self._function    = None
        self._version     = version
        self._description = description
        self._tags        = tags

    def __call__(self,):

        if not self._function:
        
            @orchestrator(
                tags=self._tags,
                description=self._description
            )
            def wrapper(*args, **kwargs):
                return function(*args, **kwargs)

            self._function = wrapper

        return self._function

    def metadata(self):
        return {
            "version": self._version,
            "description": self._description,
            "tags": self._tags,
        }

    # Other methods...

@step(
    tags = ["test"],
    description = "Example description"
    version = "0.1"
)
def test(a,b):
    return a+b

This is an improvement since:

  • The function is decorated
  • It applies the orchestrator decorator correctly
  • The function is decorated and saved, so it isn't decorated every time I call it But it doesn't allow me to access the methods, since test is the decorated test and not an instance of step.

This allows me to access the metadata with print(test.version), etc. But I can't add methods like I was trying with test.metadata().

I could get around some of the functionality by writing:

class step:

    def __init__(self, version, description, tags):

        self._function    = None
        self._version     = version
        self._description = description
        self._tags        = tags

    def __call__(self,):

        if not self._function:
            
            @orchestrator(
                tags=self._tags,
                description=self._description
            )
            def wrapper(*args, **kwargs):
                return function(*args, **kwargs)

            self._function = wrapper
            self._function.version = self._version
            self._function.description = self._description
            self._function.tags = self._tags

        return self._function

@step(
    tags = ["test"],
    description = "Example description"
    version = "0.1"
)
def test(a,b):
    return a+b

As @Jasmijn wrote, this last approach has the same functionality as:

def step(tags, description, version):
    def _inner(f):
        f = orchestrator(tags=tags, version=version)(f)
        f.tags = tags
        f.description = description
        f.version = version
        return f
return _inner

But that removes the flexibility to add methods down the line (exemplified by test.metadata(), but more complex functionality may be needed).

When using a class as a decorator that takes arguments, what is the pythonic way to get the instance of the class instead of the .__call__() method?


Solution

  • Let's take a step back. You'd like to write:

    @step(
        tags = ["test"],
        description = "Example description"
        version = "0.1"
    )
    def test(a,b):
        return a+b
    

    to be equivalent to:

    @orchestrator(
        tags = ["test"],
        version = "0.1"
    )
    def test(a,b):
        return a+b
    
    test.tags = ["test"]
    test.description = "Example description"
    test.version = "0.1"
    

    Right?

    Let's forget about the class for now and just write the simplest code that could possibly work:

    def step(tags, description, version):
        def _inner(f):
            f = orchestrator(tags=tags, version=version)(f)
            f.tags = tags
            f.description = description
            f.version = version
            return f
        return _inner
    

    Added in response to comment:

    We can make test the instance of a custom class as well. It wouldn't be the decorator, though. Here's how I would approach it:

    class OrchestratedFunction:
        def __init__(self, function, tags, version, description):
            self.function = function
            self.tags = tags
            self.version = version
            self.description = description
        def __call__(self, *args, **kwargs):
            return self.function(*args, **kwargs)
        # add methods here!
    
    def step(tags, description, version):
        def _inner(f):
            return orchestrator(tags=tags, version=version)(OrchestratedFunction(f, tags, version, description))
        return _inner
    

    You could even add the class as an argument to step so you can have different classes for different instrumented functions:

    def step(tags, description, version, cls=OrchestratedFunction):
        def _inner(f):
            return orchestrator(tags=tags, version=version)(cls(f, tags, version, description))
        return _inner