Python 断言函数是在某个“with”语句的上下文中调用的

Question

在 python 中，我想检查给定函数是否在给定类型的 with 语句中被调用

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)

我可以做一些事情，比如强制将 arg 传递给 foo 并在装饰器中包装函数，即

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)

但那样的话我就得到处转bar了。

因此我想看看是否有办法访问 with 上下文

注意：这在现实世界中的应用是确保 mlflow.pyfunc.log_model is called within an mlflow.ActiveRun 上下文，否则它会使 ActiveRun 处于打开状态，稍后会导致问题

Answer 1

这是一种丑陋的方式：全局状态。

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

不幸的是，我认为没有任何 non-ugly 方法可以做到这一点。如果您不打算自己传递一个 Bar 实例，那么您必须依赖其他地方存在的某个状态来告诉您一个 Bar 实例存在并且当前正被用作上下文管理器。

避免全局状态的唯一方法是将状态存储在实例中，这意味着装饰器需要是实例方法并且实例需要在声明函数之前存在：

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)

这仍然是“全局状态”，因为函数 foo 现在持有对持有状态的 Bar 实例的引用。但如果 foo 永远只是一个本地函数，它可能会更容易接受。

Python 断言函数是在某个“with”语句的上下文中调用的

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

python

mlflow