Wrap all methods of a class to suppress an error in Python

Categories:

Problem

Let there be an object with a lot of methods which can possibly raise some exception.

Desired behavior of those methods when exception occurs: try to fix object and run method one more time.

Solution

Wrap all the methods using a class decorator:

 1#!/usr/bin/env python
 2import logging
 3from functools import wraps
 4from types import FunctionType
 5
 6
 7logger = logging.getLogger(__name__)
 8
 9
10class NastyError(Exception):
11    def __init__(self, arg: str) -> None:
12        self.arg = arg
13
14    def __str__(self) -> str:
15        return f'Error with arg={self.arg}'
16
17
18def nasty_function(arg):
19    raise NastyError(arg)
20
21
22def suppress_nasty_error(cls: type) -> type:
23    def suppress_and_retry(method):
24        @wraps(method)
25        def wrapper(self, *args, **kwargs):
26            try:
27                method(self, *args, **kwargs)
28            except NastyError as e:
29                logger.warning('%s', str(e))
30                ...  # FIXME: do something with the `self` object to fix the error
31                method(self, *args, **kwargs)  # Try one last time
32
33        return wrapper
34    
35    for name, attr in cls.__dict__.items():
36        if isinstance(attr, FunctionType) and not name.startswith('__'):
37            setattr(cls, name, suppress_and_retry(attr))
38
39    return cls
40
41
42@suppress_nasty_error
43class A:
44    def __init__(self, arg: str) -> None:
45        self.arg = arg
46
47    def __str__(self):
48        return f'Object A with arg={self.arg}'
49
50    def method1(self):
51        if not self.arg:
52            nasty_function(self.arg)
53        print(f'arg={self.arg}')
54
55    def method2(self, arg):
56        nasty_function(arg)
57
58
59if __name__ == '__main__':
60    a = A('')
61    a.method1()
62    a.method2('')