CODEWITHSUNDEEP
pythonpython-decorators
Jame H
Jame H

Reputation: 1384

Decorator to check input type in a class

By using decorator I can use to check variable type of function. Something like this:

def accepts(*types):
    def check_accepts(f):
        assert len(types) == f.__code__.co_argcount
        def new_f(*args, **kwds):
            for (a, t) in zip(args, types):
                assert isinstance(a, t), \
                       "arg %r does not match %s" % (a,t)
            return f(*args, **kwds)
        new_f.__name__ = f.__name__
        return new_f
    return check_accepts

@accepts(int, (int,float))
def func(arg1, arg2):
    return arg1 * arg2

func(3, 2) # -> 6
func('3', 2) # -> AssertionError: arg '3' does not match <type 'int'>

It work well for normal simple. But i don't know how use it in class:

class Test()
   @accepts(pd.DataFrame)
   def print_dataframe(df):
        print(df)

I can't initital new class and use it. Please show me how to use it

Upvotes: 0

Views: 337

Answers (1)

chepner
chepner

Reputation: 531808

As written, print_dataframe is an instance method, so you either need to take that into account:

NO_CHECK = object()

def accepts(*types):
    def check_accepts(f):
        assert len(types) == f.__code__.co_argcount
        def new_f(*args, **kwds):
            for (a, t) in zip(args, types):
                if t is NO_CHECK:
                    continue
                assert isinstance(a, t), \
                       "arg %r does not match %s" % (a,t)
            return f(*args, **kwds)
        new_f.__name__ = f.__name__
        return new_f
    return check_accepts

# Using NO_CHECK because it's easier to skip checking the self argument
# than it is to handle a check for that as-yet-undefined class Test
class Test():
   @accepts(NO_CHECK, pd.DataFrame)
   def print_dataframe(self, df):
        print(df)

or redefine print_dataframe to be a static method:

   class Test():
       @staticmethod
       @accepts(pd.DataFrame)
       def print_dataframe(df):
            print(df)

Upvotes: 1

Related Questions