Search code examples
pythonclassimportdecoratorpython-decorators

Decorator on class makes class NoneType when importing


I'm trying to understand why a decorator on my class modifies the class in such way that it appears to be 'NoneType' when trying to import the class from another script.

In my_class.py I have:

my_dict = dict()

def register(cls):
    name = cls.__name__
    my_dict[name] = cls

@register  # will be commented
class MyClass:
    def my_method(self):
        print("running class method")

print("my_dict: ", my_dict)

In another module my_main.py I import the class like

from my_class import MyClass

print(type(MyClass))
print(MyClass.my_method)

If I run it with $ python3 my_main.py I get the following output:

my_dict:  {'MyClass': <class 'my_class.MyClass'>}
<class 'NoneType'>
Traceback (most recent call last):
  File "my_main.py", line 4, in <module>
    print(MyClass.my_method)
AttributeError: 'NoneType' object has no attribute 'my_method'

By commenting the @register line in my_class.py the my_main.py runs without error and outputs:

my_dict:  {}
<class 'type'>
<function MyClass.my_method at 0x7ff06f254f28>

..but obviously my_dict is no longer filled. Is there a way of registering my_class with the given decorator AND accessing the attributes of the class after importing it in another script?


Solution

  • A decorator is no more than just a normal function.

    I don't want to describe a lot about how to write a decorator rightly. But at least, you should let your decorator return a class.

    That is similar when you write a decorator for a function. Decorator should at least return a function.

    For example:

    def decorator(method):
        @functools.wraps(method)
        def wrapper(*args, **kwargs):
            print("Hi")
            return method(*args, **kwargs)
        return wrapper
    

    When you use this decorator to decorate a function, you are actually pass that function to decorator: new_function = decorator(original_function).

    Which means new_function is wrapped by wrapper. That is how decorator works. When you execute decorated function, it actually execute the wrapper:

    print("Hi")
    return method(*args, **kwargs)  # pass all args to original function and return its return value.
    

    In your code, your decorator just returns None.