How to wrap a class's methods?

[Jeff Shannon]

Steven Bethard wrote:

> Grant Edwards wrote:
> 
>> I want to subclass an IMAP connection so that most of the
>> methods raise an exception if the returned status isn't 'OK'.
>> This works, but there's got to be a way to do it that doesn't
>> involve so much duplication:
>>
>> class MyImap4_ssl(imaplib.IMAP4_SSL):
>>
>>     def login(*args):
>>         s,r = imaplib.IMAP4_SSL.login(*args)
>>         if s!='OK':
>>             raise NotOK((s,r))
>>         return r
>>             def list(*args):
>>         s,r = imaplib.IMAP4_SSL.list(*args)
>>         if s!='OK':
>>             raise NotOK((s,r))
>>         return r
>>             def search(*args):
>>         s,r = imaplib.IMAP4_SSL.search(*args)
>>         if s!='OK':
>>             raise NotOK((s,r))
>>         return r
>>             [and so on for another dozen methods]
> 
> 
> You could try something like (Untested!):
> 
> class Wrapper(object):
>     def __init__(self, func):
>         self.func = func
>     def __call__(*args, **kwargs):
>         self, args = args[0], args[1:]
>         s, r = self.func(*args)
>         if s != 'OK':
>             raise NotOK((s, r))
>         return r
> 
> for func_name in ['login', 'list', 'search']:
>     func = Wrapper(getattr(imaplib.IMAP4_SSL, func_name))
>     setattr(imaplib.IMAP4_SSL, func_name, func)

You could probably also do this as a factory function, rather than as 
a class (also untested!):

def Wrapper(func):
     def wrapped(self, *args, **kwargs):
         s, r = func(self, *args, **kwargs)
         if s != 'OK':
             raise NotOK((s,r))
         return r
     return wrapped

I believe that this will be semantically almost equivalent, but 
conceptually slightly simpler.

Jeff Shannon