How to wrap a class's methods?

[Steven Bethard]

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)

STeVe