diff --git a/rainbowadn/instrument/instrumentation.py b/rainbowadn/instrument/instrumentation.py index fe65171..6bf3ba8 100644 --- a/rainbowadn/instrument/instrumentation.py +++ b/rainbowadn/instrument/instrumentation.py @@ -1,13 +1,13 @@ import functools from contextlib import ExitStack -from typing import Callable, Generic, TypeVar +from typing import Callable, TypeVar __all__ = ('Instrumentation',) -IType = TypeVar('IType') +IType = TypeVar('IType', bound='Instrumentation') -class Instrumentation(Generic[IType]): +class Instrumentation: deinstrumentation = {} _method: Callable _wrap: Callable @@ -24,8 +24,6 @@ class Instrumentation(Generic[IType]): raise NotImplementedError def __enter__(self: IType) -> IType: - if not isinstance(self, Instrumentation): - raise TypeError if hasattr(self, '_method') or hasattr(self, '_wrap'): raise RuntimeError method = getattr(self.target, self.methodname)