How to wrap each method in a class with a context manager?

Question

I have a class Class, that has a certain number of methods. I want to wrap each of those methods with a context manager, such that when calling Class().foo, it will run Class.foo in the context. Here's the code I have so far:

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)))

I'm planning on calling this function (but defined before) at the end of a class definition, so after all the methods are defined

I think what you need here is a metaclass - that will allow you to modify how the class definition is taken and turned into a class i.e. you can wrap all the methods in context manager decorator see https://stackoverflow.com/questions/100003/what-are-metaclasses-in-python and https://www.pythontutorial.net/python-oop/python-metaclass/ — Anentropic, Jun 16 '22 at 15:56
@Anentropic I looked at it and [this](https://stackoverflow.com/a/11350487/17920058), but it looked far too complicated. If there's an easier way, I would like to know. I did find a way to fix what I have so it can be used. — DrownedSuccess, Jun 16 '22 at 16:35
I dunno, metaclass seems like the right way to me. The answers and tutorials for metaclasses are all complicated to read but if you can get past that the concept itself is not too complicated. I can't tell from the code posted in the question how you actually use what you have — Anentropic, Jun 16 '22 at 16:50
@Anentropic I put my modified function my answer (I did not know you can reference other methods in `__init__`, which makes it much easier for me). — DrownedSuccess, Jun 16 '22 at 16:51
@Anentropic Added an example (the actual class I'm using is far more complicated). — DrownedSuccess, Jun 16 '22 at 16:54

score 3 · Accepted Answer · answered Jun 16 '22 at 17:28

Here is a metaclass version:

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())

You could also do away with the context manager and have the same thing just in the decorator definition, i.e.:

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

One advantage of the metaclass version is that the work to decorate the methods of the class happens only once, at "import time", instead of each time you make an instance. So that would be a little more efficient if you are making lots of instances and there are many methods to decorate.

While my answer worked for me, this answer is probably much more Pythonic and sustainable. — DrownedSuccess, Jun 16 '22 at 20:02

DrownedSuccess · Answer 2 · 2022-06-16T17:33:23.527

This works:

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)))

You just need to call it in your class's __init__ method.

Usage:

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())

How to wrap each method in a class with a context manager?

2 Answers2