Search code examples
pythondecoratormetaclass

Using decorators to eliminate need for if statements in methods


I have a class which stores some data from a file in it, but sometimes the file don't exist.

The class provides multiple methods to do calculations on data, but all those methods should raise error if no data is found.

This is an example of class and one of those methods:

class A:
    def __init__(self, file_name):
        if isfile(file_name):
            with open(file_name) as f:
                self.data = json.load(f)
        else:
            self.data = None

   def get_something(self):
       if self.data is None:
          raise ValueError('no data')

       return data['a'] + data['b']

Question How can I get rid of having an if statement in every class method. Is there a way to use decorators on the class or method level, which makes it look better than scattered if statements?


Solution

  • If you want to have a decorator at the method level, define a wrapper function and look at the self argument to figure out what data is:

    def requires_data(a_method):
        def wrapped_method(self, *args, **kwargs):
            assert isinstance(self, A), "only use this decorator on A methods"
            if self.data is None:
                raise ValueError('no data')
            return a_method(self, *args, **kwargs)
        return wrapped_method
    
    
    class A:
        def __init__(self, data):
            self.data = data
    
        @requires_data
        def get_something(self):
            return self.data['a'] + self.data['b']
    
    
    print(A({'a': 21, 'b': 21}).get_something())
    print(A(None).get_something())  # raises!
    

    As noted in the comments, though, I would instead favor having data be a required attribute, such that __init__ raises immediately if it's not available:

    class A:
        def __init__(self, data):
            if data is None:
                raise ValueError('no data')
            self.data = data
    
        def get_something(self):
            return self.data['a'] + self.data['b']
    
    
    print(A({'a': 21, 'b': 21}).get_something())
    print(A(None).get_something())  # raises!
    

    If the class has some functionality that depends on data and some that doesn't, and there are some cases where you want only the non-data functionality, split the class apart along those lines.