Python assert a function is called within a certain `with` statement's context-CodePudding

In python I would like to check that a given function is called within a with statement of a given type

class Bar:
 def __init__(self, x):
  self.x = x
 def __enter__(self):
  return self
 def __exit__(self, *a, **k):
  pass

def foo(x):
 # assert that the enclosing context is an instance of bar
 # assert isinstance('enclosed context', Bar)
 print(x*2)

with Bar(1) as bar:
 foo(bar.x)

I could do something like enforcing an arg passed into foo and wrapping functions in a decorator i.e.

class Bar:
 def __init__(self, x):
  self.x = x
 def __enter__(self):
  return self
 def __exit__(self, *a, **k):
  pass

def assert_bar(func):
 def inner(bar, *a, **k):
  assert isinstance(bar, Bar)
  return func(*a, **k)
 return inner


@assert_bar
def foo(x):
 print(x*2)

with Bar(1) as bar:
 foo(bar, bar.x)

but then I would have to pass around bar everywhere.

As a result I'm trying to see if there's a way to access the with context

Note: The real world application of this is ensuring that mlflow.pyfunc.log_model is called within an mlflow.ActiveRun context, or it leaves an ActiveRun open, causing problems later on

CodePudding user response：

Here's an ugly way to do it: global state.

class Bar:
    active = 0
    def __init__(self, x):
        self.x = x
    def __enter__(self):
        Bar.active  = 1
        return self
    def __exit__(self, *a, **k):
        Bar.active -= 1

from functools import wraps

def assert_bar(func):
    @wraps(func)
    def wrapped(*vargs, **kwargs):
        if Bar.active <= 0:
            # raises even if asserts are disabled
            raise AssertionError()
        return func(*vargs, **kwargs)
    return wrapped

Unfortunately I don't think there is any non-ugly way to do it. If you aren't going to pass around a Bar instance yourself then you must rely on some state existing somewhere else to tell you that a Bar instance exists and is currently being used as a context manager.

The only way you can avoid that global state is to store the state in the instance, which means the decorator needs to be an instance method and the instance needs to exist before the function is declared:

from functools import wraps

class Bar:
    def __init__(self, x):
        self.x = x
        self.active = 0
    def __enter__(self):
        self.active  = 1
        return self
    def __exit__(self, *a, **k):
        self.active -= 1
    def assert_this(self, func):
        @wraps(func)
        def wrapped(*vargs, **kwargs):
            if self.active <= 0:
                raise AssertionError()
            return func(*vargs, **kwargs)
        return wrapped

bar = Bar(1)

@bar.assert_this
def foo(x):
    print(x   1)

with bar:
    foo(1)

This is still "global state" in the sense that the function foo now holds a reference to the Bar instance that holds the state. But it may be more palatable if foo is only ever going to be a local function.