如何用上下文管理器包装类中的每个方法?

Question

我有一个类Class，它有一定数量的方法。我想用上下文管理器包装这些方法，这样当调用Class().foo时，它将在上下文中运行Class.foo。下面是我到目前为止的代码:

def wrap_all_methods_in_class_with_chdir_contextmanager(self):
@contextlib.contextmanager
def set_directory(path):
"""Sets the cwd within the context

Args:
path (Path): The path to the cwd

Yields:
None
"""

origin = os.path.abspath(os.getcwd())
try:
os.chdir(path)
yield
finally:
os.chdir(origin)

def wrapper(func):
def new_func(*args, **kwargs):
with set_directory(f"{ROOT}/{self.name}"):
getattr(self,func)(*args, **kwargs)
return new_func

for func in [func for func in dir(self) if callable(getattr(self, func)) and not func.startswith('__')]:
setattr(self,func,wrapper(getattr(self,func)))

我计划在类定义结束时调用这个函数(但之前定义过)，所以在定义了所有方法之后

Answer 1

这是一个元类版本:

import contextlib
import os
from functools import wraps
from types import FunctionType
ROOT = "/base"
@contextlib.contextmanager
def set_directory(path):
"""
Sets the cwd within the context
Args:
path (Path): The path to the cwd
Yields:
None
"""
origin = os.path.abspath(os.getcwd())
try:
os.chdir(path)
yield
finally:
os.chdir(origin)
def use_chdir_context(f):
@wraps(f)
def wrapped(self, *args, **kwargs):
with set_directory(f"{ROOT}/{self.path}"):
return f(self, *args, **kwargs)
return wrapped

class ChdirContextMetaclass(type):
"""
Wraps all methods in the class with use_chdir_context decorator
"""
def __new__(metacls, name, bases, class_dict):
new_class_dict = {}
for attr_name, attr in class_dict.items():
if isinstance(attr, FunctionType):
attr = use_chdir_context(attr)
new_class_dict[attr_name] = attr
return type.__new__(metacls, name, bases, new_class_dict)

class Hi(metaclass=ChdirContextMetaclass):
path = "hi/1/2/3"
def lovely(self):
print(os.getcwd())

你也可以去掉上下文管理器，只在装饰器定义中做同样的事情，即:

def use_chdir_context(f):
@wraps(f)
def wrapped(self, *args, **kwargs):
origin = os.path.abspath(os.getcwd())
try:
os.chdir(f"{ROOT}/{self.path}")
return f(self, *args, **kwargs)
finally:
os.chdir(origin)
return wrapped

元类版本的一个优点是，修饰类方法的工作只发生一次，在"导入时"，而不是每次创建实例时。如果你创建了很多实例并且有很多方法可以装饰的话这会更有效率。

Answer 2

这行得通:

def wrap_all_methods_in_class_with_chdir_contextmanager(self,path):
@contextlib.contextmanager
def set_directory(path):
"""Sets the cwd within the context

Args:
path (Path): The path to the cwd

Yields:
None
"""

origin = os.path.abspath(os.getcwd())
try:
os.chdir(path)
yield
finally:
os.chdir(origin)

def wrapper(func):
def new_func(*args, **kwargs):
with set_directory(path):
return func(*args, **kwargs)
return new_func

for func in [func for func in dir(self) if callable(getattr(self, func)) and not func.startswith('__')]:
setattr(self,func,wrapper(getattr(self,func)))

你只需要在类的__init__方法中调用它。

使用:

class hi:
def __init__(self):
wrap_all_methods_in_class_with_chdir_contextmanager(self,"/tmp")
def lovely(self):
print(os.getcwd())

def foo(self):
print(os.getcwd())
test=hi()
test.foo()
print(os.getcwd())