how to use python decorator with argument?

Question

I would like to define a decorator that will register classes by a name given as an argument of my decorator. I could read from stackoverflow and other sources many examples that show how to derive such (tricky) code but when adapted to my needs my code fails to produce the expected result. Here is the code:

import functools

READERS = {}

def register(typ):
    def decorator_register(kls):
        @functools.wraps(kls)
        def wrapper_register(*args, **kwargs):
            READERS[typ] = kls 
        return wrapper_register
    return decorator_register

@register(".pdb")
class PDBReader:
    pass

@register(".gro")
class GromacsReader:
    pass

print(READERS)

This code produces an empty dictionary while I would expect a dictionary with two entries. Would you have any idea about what is wrong with my code?

Answer 1

Taking arguments (via (...) ) and decoration (via @ ) both result in calls of functions. Each "stage" of taking arguments or decoration maps to one call and thus one nested functions in the decorator definition. register is a three-stage decorator and takes as many calls to trigger its innermost code. Of these,

• the first is the argument ( (".pdb") ),
• the second is the class definition ( @... class ), and
• the third is the class call/instantiation ( PDBReader(...) )
  • This stage is broken as it does not instantiate the class.

In order to store the class itself in the dictionary, store it at the second stage. As the instances are not to be stored, remove the third stage.

def register(typ):                # first stage: file extension
    """Create a decorator to register its target for the given `typ`"""
    def decorator_register(kls):  # second stage: Reader class
        """Decorator to register its target `kls` for the previously given `typ`"""
        READERS[typ] = kls 
        return kls                # <<< return class to preserve it
    return decorator_register

Take note that the result of a decorator replaces its target. Thus, you should generally return the target itself or an equivalent object. Since in this case the class is returned immediately, there is no need to use functools.wraps .

READERS = {}

def register(typ):                # first stage: file extension
    """Create a decorator to register its target for the given `typ`"""
    def decorator_register(kls):  # second stage: Reader class
        """Decorator to register its target `kls` for the previously given `typ`"""
        READERS[typ] = kls 
        return kls                # <<< return class to preserve it
    return decorator_register

@register(".pdb")
class PDBReader:
    pass

@register(".gro")
class GromacsReader:
    pass

print(READERS)  # {'.pdb': <class '__main__.PDBReader'>, '.gro': <class '__main__.GromacsReader'>}

Answer 2

If you don't actually call the code that the decorator is "wrapping" then the "inner" function will not fire, and you will not create an entry inside of READER . However, even if you create instances of PDBReader or GromacsReader , the value inside of READER will be of the classes themselves, not an instance of them.

If you want to do the latter, you have to change wrapper_register to something like this:

def register(typ):
    def decorator_register(kls):
        @functools.wraps(kls)
        def wrapper_register(*args, **kwargs):
            READERS[typ] = kls(*args, **kwargs)
            return READERS[typ]
        return wrapper_register
    return decorator_register

I added simple init/repr inside of the classes to visualize it better:

@register(".pdb")
class PDBReader:
    def __init__(self, var):
        self.var = var
    def __repr__(self):
        return f"PDBReader({self.var})"

@register(".gro")
class GromacsReader:
    def __init__(self, var):
        self.var = var
    def __repr__(self):
        return f"GromacsReader({self.var})"

And then we initialize some objects:

x = PDBReader("Inside of PDB")
z = GromacsReader("Inside of Gromacs")
print(x) # Output: PDBReader(Inside of PDB)
print(z) # Output: GromacsReader(Inside of Gromacs)
print(READERS) # Output: {'.pdb': PDBReader(Inside of PDB), '.gro': GromacsReader(Inside of Gromacs)}

---

If you don't want to store the initialized object in READER however, you will still need to return an initialized object, otherwise when you try to initialize the object, it will return None .

You can then simply change wrapper_register to:

def wrapper_register(*args, **kwargs):
    READERS[typ] = kls
    return kls(*args, **kwargs)