diff --git a/rainbowadn/instrument/instrumentation.py b/rainbowadn/instrument/instrumentation.py index 18147c5..bb8be28 100644 --- a/rainbowadn/instrument/instrumentation.py +++ b/rainbowadn/instrument/instrumentation.py @@ -58,3 +58,10 @@ class Instrumentation(Generic[IType]): def enter(self: IType, stack: ExitStack) -> IType: assert isinstance(self, Instrumentation) return stack.enter_context(self) + + def enter_conditional(self: IType, stack: ExitStack) -> IType | None: + assert isinstance(self, Instrumentation) + if hasattr(self.target, self.methodname): + return self.enter(stack) + else: + return None