Source code for datumaro.util.scope

# Copyright (C) 2021 Intel Corporation
#
# SPDX-License-Identifier: MIT

from contextlib import ExitStack, contextmanager
from functools import partial, wraps
from typing import Any, Callable, ContextManager, Dict, Optional, Tuple, TypeVar
import threading

from attr import attrs

from datumaro.util import optional_arg_decorator

T = TypeVar('T')

[docs]class Scope: """ A context manager that allows to register error and exit callbacks. """ _thread_locals = threading.local()
[docs] @attrs(auto_attribs=True) class ExitHandler: callback: Callable[[], Any] ignore_errors: bool = True def __exit__(self, exc_type, exc_value, exc_traceback): try: self.callback() except Exception: if not self.ignore_errors: raise
[docs] @attrs class ErrorHandler(ExitHandler): def __exit__(self, exc_type, exc_value, exc_traceback): if exc_type: return super().__exit__(exc_type=exc_type, exc_value=exc_value, exc_traceback=exc_traceback)
[docs] def __init__(self): self._stack = ExitStack() self.enabled = True
[docs] def on_error_do(self, callback: Callable, *args, kwargs: Optional[Dict[str, Any]] = None, ignore_errors: bool = False): """ Registers a function to be called on scope exit because of an error. If ignore_errors is True, the errors from this function call will be ignored. """ self._register_callback(self.ErrorHandler, ignore_errors=ignore_errors, callback=callback, args=args, kwargs=kwargs)
[docs] def on_exit_do(self, callback: Callable, *args, kwargs: Optional[Dict[str, Any]] = None, ignore_errors: bool = False): """ Registers a function to be called on scope exit. """ self._register_callback(self.ExitHandler, ignore_errors=ignore_errors, callback=callback, args=args, kwargs=kwargs)
def _register_callback(self, handler_type, callback: Callable, args: Tuple[Any] = None, kwargs: Dict[str, Any] = None, ignore_errors: bool = False): if args or kwargs: callback = partial(callback, *args, **(kwargs or {})) self._stack.push(handler_type(callback, ignore_errors=ignore_errors))
[docs] def add(self, cm: ContextManager[T]) -> T: """ Enters a context manager and adds it to the exit stack. Returns: cm.__enter__() result """ return self._stack.enter_context(cm)
[docs] def enable(self): self.enabled = True
[docs] def disable(self): self.enabled = False
[docs] def close(self): self.__exit__(None, None, None)
def __enter__(self) -> 'Scope': return self def __exit__(self, exc_type, exc_value, exc_traceback): if not self.enabled: return self._stack.__exit__(exc_type, exc_value, exc_traceback) self._stack.pop_all() # prevent issues on repetitive calls
[docs] @classmethod def current(cls) -> 'Scope': return cls._thread_locals.current
[docs] @contextmanager def as_current(self): previous = getattr(self._thread_locals, 'current', None) self._thread_locals.current = self try: yield finally: self._thread_locals.current = previous
[docs]@optional_arg_decorator def scoped(func, arg_name=None): """ A function decorator, which allows to do actions with the current scope, such as registering error and exit callbacks and context managers. """ @wraps(func) def wrapped_func(*args, **kwargs): with Scope() as scope: if arg_name is None: with scope.as_current(): ret_val = func(*args, **kwargs) else: kwargs[arg_name] = scope ret_val = func(*args, **kwargs) return ret_val return wrapped_func
# Shorthands for common cases
[docs]def on_error_do(callback, *args, ignore_errors=False, kwargs=None): return Scope.current().on_error_do(callback, *args, ignore_errors=ignore_errors, kwargs=kwargs)
[docs]def on_exit_do(callback, *args, ignore_errors=False, kwargs=None): return Scope.current().on_exit_do(callback, *args, ignore_errors=ignore_errors, kwargs=kwargs)
[docs]def scope_add(cm: ContextManager[T]) -> T: return Scope.current().add(cm)