Anshul
Anshul

Reputation: 776

Wrapping a constructor in decorator in python

I want to have a handy decorator that checks if the value of attribute passed to the method is not None. I then want to use it as a general purpose decorator in class methods. For this, I wrote:

def check_empty(name):
  def wrap_outer(function):
    def wrapped(*args, **kwargs):
      if kwargs.get(name, None) is None:
        raise Exception('{item} cannot be empty'.format(item=name))
      return function(args, kwargs)
    return wrapped
  return wrap_outer

And, the class is defined as:

class Player(object):
  @check_empty('name')
  def __init__(self, name):
    self.name = name

def __str__(self):
  return self.name

However, this does not work. The template for the constructor and the decorator do not match. How can we create a function decorator for a constructor? Or is class based decorator a better option?


Thanks to @ajax1234, I figured out the problem. It was in the way I was passing down my args and kwargs. It should have been:

def check_empty(attr):
  def wrap_outer(function):
    def wrapped(*args, **kwargs):
      if (len(args) > 1 and args[1] is None) or (len(args)==1 and len(kwargs)==0):
        raise Exception('{item} cannot be empty'.format(item=attr))
      elif attr in kwargs and kwargs.get(attr, None) is None:
        raise Exception('{item} cannot be empty'.format(item=attr))
    return function(*args, **kwargs)
  return wrapped
return wrap_outer

This takes care of checking on the parameter if it is passed as a keyword argument or just a positional argument.

NOTE: It does have a restriction that positional argument logic won't scale well while maintaining the generic nature. (We'd need to look into the positional argument's position using inspect.getargspec and then build up from there.)

Upvotes: 0

Views: 5024

Answers (2)

Ryan Stein
Ryan Stein

Reputation: 8000

The only way I can see you being able to do what you want, while including positional and keyword arguments, is to provide the mapping of positions to names, like so:

class NoneException(Exception):
    pass


def check_none(**check):
    def wrapper(fn):
        def wrapped(*args, **kwargs):
            for key, pos in check.items():
                try:
                    if kwargs[key] == None:
                        raise NoneException(f'{key} cannot be empty')
                except KeyError:
                    pass

                try:
                    if args[pos] == None:
                        raise NoneException(f'{key} cannot be empty')
                except IndexError:
                    pass

            return fn(*args, **kwargs)

        return wrapped

    return wrapper


class A:
    @check_none(name=1)
    def __init__(self, name):
        self.name = name

    @check_none(phrase=1)
    def test(self, phrase):
        print(f'My name is {self.name}. {phrase}!')


a = A('A')

a.test('Hello world!')

try:
    a.test(None)
except NoneException:
    print('successfully failed')
else:
    print('something went wrong')

This should cover the cases of where you provide an argument positionally and explicitly named, however note that it is a brittle solution considering that the mapping of positional argument indexes must be maintained.

If there were some way to reflectively divine the mapping self=0, name=1, et cetera from Python's internals, I would suggest to use that instead.


Alternatively, a simpler solution is a function call at the top of each function you'd want to check.

class NoneException(Exception):
    pass

def check_none(*args, **kwargs):
    for key, value in (*enumerate(args), *kwargs.items()):
        if value == None:
            raise NoneException(f'key {key} == None')

class A:
    def __init__(self, name):
        check_none(name)

        self.name = name

Upvotes: 1

Ajax1234
Ajax1234

Reputation: 71471

You can try this:

def check_empty(attr):
   def check_method(function):
     def wrapped(cls, **kwargs):
       if kwargs.get(attr) is None:
         raise Exception('{item} cannot be empty'.format(item=attr))
       return function(cls, **kwargs)
     return wrapped
   return check_method

class Player(object):
  @check_empty('name')
  def __init__(self, **kwargs):
    self.__dict__ = kwargs
  def __str__(self):
     return self.name

p = Player(name=None)

Output:

Exception: Name cannot be empty

However, any parameter other than None works:

p = Player(name = 'Foo')
print(p.name)

Output:

'Foo'

Upvotes: 3

Related Questions