Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

decorate a python class such that most methods raise an exception if condition

I have had a situation arise where most methods of a class need to raise an exception if called except for one, if a certain condition is False. It would be possible to go into most methods and write if not condition such that each method will raise the exception if the condition is not true, but I think it is probably possible to do this somehow with a single decorator on the top of the class.

This question is similar but it involves decorating every method separately, and if I were going to do that, I might as well just put the if statement into each method.

Here is some code and comments to help communicate it:

CONDITION = True  # change to False to test

def CheckMethods():
    if CONDITION:
        # run all the methods as usual if they are called
        pass 
    else:
        # raise an exception for any method which is called except 'cow'
        # if 'cow' method is called, run it as usual
        pass


@CheckMethods
class AnimalCalls:
    def dog(self):
        print("woof")
    def cat(self):
        print("miaow")
    def cow(self):
        print("moo")
    def sheep(self)
        print("baa") 

a = AnimalCalls()
a.dog()
a.cat()
a.cow()
a.sheep()

Does anyone know how to do this? Have never decorated a class before or tried to check its methods like this.

like image 397
cardamom Avatar asked Jan 01 '23 19:01

cardamom


2 Answers

Implementing a proxy is as simple as that

class Proxy:
    def __init__(self, inst):
        self.__inst = inst

    def __getattr__(self, name):
        return getattr(self.__inst, name)

Instead of obj = SomeClass() you'd use obj = Proxy(SomeClass()). All accesses to obj.attribute get intercepted by Proxy.__getattr__. That's the method you may add more logic to, e.g.:

class MethodChecker:
    def __init__(self, inst, check):
        self.__inst = inst
        self.__check = check

    def __getattr__(self, name):
        self.__check()
        return getattr(self.__inst, name)
like image 157
kay Avatar answered Jan 30 '23 22:01

kay


The proxy would be my pick, but here is a decorator as requested.

I added a test to exclude any methods starting with an underscore. You might want to include _internal methods, but take care not to mess with any special __dunder__ methods.

# cond = lambda attr: True  # full access
cond = lambda attr: attr == 'cow'

def methodcheck(cls):
    def cond_getattribute(self, name):
        if name.startswith('_') or cond(name):
            return saved_gettattribute(self, name)
        raise AttributeError("access forbidden")
    saved_gettattribute = cls.__getattribute__
    cls.__getattribute__ = cond_getattribute
    return cls 

@methodcheck
class AnimalCalls:
    def dog(self):
        print("woof")
    def cat(self):
        print("miaow")
    def cow(self):
        print("moo")
    def sheep(self):
        print("baa"
like image 41
VPfB Avatar answered Jan 30 '23 22:01

VPfB