ContextManager + access the decorated function in __enter__

[pmav99]

I have a context manager that creates some resources and cleans them up on exit. I would like, if possible, to use it as a decorator, too. But when I do, I also need to have access to the name of the decorated function. I managed to do this by subclassing decorator.ContextManager and overriding __call__:

    def __call__(self, func):
        """Context manager decorator"""
        self.__func__ = func    # this is the line I added
        return decorator.FunctionMaker.create(
            func, "with _self_: return _func_(%(shortsignature)s)",
            dict(_self_=self, _func_=func), __wrapped__=func)

1. Is it possible to achieve the same thing with "vanilla" decorator.ContextManager?
2. If not, would you accept a PR for this?

[micheles]

Doing self.__func__ = func is dangerous. If the same decorator/context manager is used to decorate two different functions (say func1 and func2) at the second time self.__func__ will be equal to func2 and the reference to func1 will be lost. Is that what you want?

[pmav99]

Of course, you are right... it does work if you use a different decorator instance but not if you reuse the same decorator.

Just for reference, full example:

import contextlib
import pathlib
import uuid

import decorator

class ContextManager(contextlib._GeneratorContextManager):
    def __init__(self, g, *a, **k):
        return contextlib._GeneratorContextManager.__init__(self, g, a, k)

    def __call__(self, func):
        self.__func__ = func
        return decorator.FunctionMaker.create(
            func,
            "with _self_: return _func_(%(shortsignature)s)",
            dict(_self_=self, _func_=func),
            __wrapped__=func,
        )

class CMFactory(ContextManager):
    def __init__(self, default_name=None):
        self.default_name = default_name or uuid.uuid4().hex

    @property
    def name(self):
        return getattr(self, "__func__", self.default_name)

    def __enter__(self):
        print(f"Entering {self.name}")

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        print(f"Exiting  {self.name}")
        print()

This works:

@CMFactory()
def func1():
    print("inside func")

@CMFactory()
def func2():
    print("inside func")

with CMFactory():
    print("Inside context")

func1()
func2()

This doesn't

cm = CMFactory()

@cm
def func1():
    print("inside func")

@cm
def func2():
    print("inside func")

with cm:
    print("Inside context")

func1()
func2()

The funny thing is that this works, too, but for the wrong reasons):

cm = CMFactory()

@cm
def func1():
    print("inside func")

func1()

@cm
def func2():
    print("inside func")

func2()

I guess, someone could prevent re-using the decorator, but it starts to feel rather hackish + it reduces the scope of the ContextManager class...:

    def __call__(self, func):
        if hasattr(self, "__func__"):
            raise ValueError("You can't reuse this decorator, please create a new instance")
        self.__func__ = func
        return decorator.FunctionMaker.create(
            func,
            "with _self_: return _func_(%(shortsignature)s)",
            dict(_self_=self, _func_=func),
            __wrapped__=func,
        )

Thank you Michele. If there are no workarounds, feel free to close

[micheles]

There are solutions, for instance to save in the context manager a dictionary func.__qualname_ -> func, but you can do that yourself, I am not convinced I should change the decorator module. It does not look such a common use case to me: why do you need to keep references to the original functions in the decorator?

[pmav99]

I am not convinced I should change the decorator module

From what I've seen so far, I would argue that you shouldn't)

why do you need to keep references to the original functions in the decorator?

Short answer: I have a context manager that I want to also use as a decorator. The context manager creates some resources in a directory that needs to have a unique name. When it is used as a decorator though, it would make sense to use the decorated function's name to make it easy to associate the directory with the function (e.g. using f"{func.__module__}_{func.__qualname__}"). Each function will only be used once.

Longer answer: I am refactoring some tests that use a custom extension framework to unittest. Unfortunately, the extension's test runner does some weird stuff. It not only runs the tests but it also sets up/tears down a "session" for each test. In other words some of the test fixtures are not part of the tests/test suite but of the test runner... We want to ditch that test runner and run the tests with pytest. But we need to implement setting up/tearing down this session in a way that will not modify the existing tests and that will let us create a session that can be directly associated to each test.

class CustomTestCase(unittest.TestCase):
    pass

class FooTestCase(CustomTestCase):
    def test1(self): 
        pass

    def test2(self): 
        pass

So, I need to create a "session" for each test that among other things, will create a temp directory that will get removed after the test finishes. Ideally, I would also like to be able to optionally keep the directory for inspection (e.g. when the tests are failing). Now, I already have a context manager that implements this. E.g

with temp_session(name='my_temp_session", cleanup=True):
    # This context manager creates a directory named "my_temp_session" which gets removed on exit
    # if `cleanup=False` the directory will not get removed.

So I wanted to reuse this as a decorator that would do the same but which instead of a random name for the temp directory would use the name of the decorated function. This would make it easier to associate the directory with each failing test.

# contents of tests/test_foo.py

class FooTestCase(unittest.TestCase):
    @temp_session() # this would create a directory named "tests_test_foo_Foo_TestCase_test1"
    def test1(self): 
        pass

    @temp_session(cleanup=False) # this would create a directory named "tests_test_foo_Foo_TestCase_test2"
    def test2(self): 
        pass

Of course I could implement both a context manager and a decorator and move on - which is what I am probably going to do. Or just use random names and be done with it. But the idea of using the same class both for the context manager and the decorator was tempting)

There are solutions, for instance to save in the context manager a dictionary func._qualname -> func

I am probably missing something obvious here - I also thought about creating a registry of some sort - but I think that the problem remains; i.e. you can't access func from within __enter__ so you can't retrieve the "correct" value from the dictionary, am I wrong?

[micheles]

I also have tests that require creating a temporary directory. What I did was to introduce a custom subclass of unittest.TestCase with a redefined run method taking care of the directories. I did not use decorators at all. Something like this:

import os
import tempfile
import unittest
import shutil

tmp = tempfile.gettempdir()  # /tmp, but it would be preferable to create a new directory

class BaseTestCase(unittest.TestCase):
    def run(self, result=None):
        dtemp = os.path.join(tmp, self._testMethodName)
        if not os.path.exists(dtemp):
            os.mkdir(dtemp)
            print('creating', dtemp)
        result = super().run()
        shutil.rmtree(dtemp)
        print('removing', dtemp)
        return result


class ExampleTestCase(BaseTestCase):
    def test1(self):
        pass

    def test2(self):
        pass

that you can run with pytest:

$ pytest -s x.py
============================= test session starts ==============================
platform linux -- Python 3.6.7, pytest-4.2.0, py-1.7.0, pluggy-0.8.1
rootdir: /home/michele, inifile:
plugins: xdist-1.27.0, forked-1.0.2, celery-4.1.0
collected 2 items                                                              

x.py creating /tmp/test1
removing /tmp/test1
.creating /tmp/test2
removing /tmp/test2
.