1def optimize_exception_check(attr_name):
 2    def decorator(func):
 3        def wrapper(self, exc_type):
 4            cache = getattr(self, attr_name)
 5            if exc_type in cache:
 6                return True
 7
 8            result = func(self, exc_type)  # True or False
 9            if result:
10                cache.add(exc_type)
11            return result
12        return wrapper
13    return decorator
14
15
16class ProtectedSection:
17    __counter = 0
18
19    def __init__(self, log=(), suppress=()):
20        if log is not None:
21            for exc in log:
22                if not issubclass(exc, Exception):
23                    raise TypeError("Not a valid exception type")
24
25        if suppress is not None:
26            for exc in suppress:
27                if not issubclass(exc, Exception):
28                    raise TypeError("Not a valid exception type")
29
30        self._current_exception = None  # Saves the specific or last "updated" exception in case of reuse of an already created instance
31        self.log_exceptions = set(log)
32        self.suppress_exceptions = set(suppress)
33        self._checked_exceptions_logs = set()  # Optimization to avoid checking one exception n times
34        self._checked_exceptions_suppress = set()  # Optimization to avoid checking one exception n times
35        self.exceptions_by_session = {}  # Dictionary to store exceptions with their session ID
36
37    @optimize_exception_check("_checked_exceptions_logs")
38    def is_exception_logged(self, exc_type):
39        return exc_type in self.log_exceptions
40
41    @optimize_exception_check("_checked_exceptions_suppress")
42    def is_exception_suppressed(self, exc_type):
43        return exc_type in self.suppress_exceptions
44
45    def __enter__(self):
46        self._current_exception = None
47        return self
48
49    def __exit__(self, exc_type, exc_value, traceback):
50        if exc_type:
51            if self.is_exception_logged(exc_type):
52                self.exceptions_by_session[self.__counter] = exc_value
53                self.__counter += 1
54                self._current_exception = exc_value
55                return True
56
57            if self.is_exception_suppressed(exc_type):
58                self._current_exception = None
59                self.__counter += 1
60                return True
61
62        return False
63
64    @property
65    def exception(self):
66        return self._current_exception
67
68    @property
69    def exceptions(self):
70        return self.exceptions_by_session
71
72    def get_exception_by_session(self, session_id):
73        return self.exceptions_by_session.get(session_id, None)