How to create a python class with a single use context

Question

If we look at python docs it states:

Most context managers are written in a way that means they can only be used effectively in a with statement once. These single use context managers must be created afresh each time they’re used - attempting to use them a second time will trigger an exception or otherwise not work correctly.

This common limitation means that it is generally advisable to create context managers directly in the header of the with statement where they are used (as shown in all of the usage examples above).

Yet, the example most commonly shared for creating context managers inside classes is:

from contextlib import ContextDecorator
import logging

logging.basicConfig(level=logging.INFO)

class track_entry_and_exit(ContextDecorator):
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        logging.info('Entering: %s', self.name)

    def __exit__(self, exc_type, exc, exc_tb):
        logging.info('Exiting: %s', self.name)

But, when I instantiate this class, I can pass it several times to a with statement:

In [8]: test_context = track_entry_and_exit('test')

In [9]: with test_context:
   ...:     pass
   ...: 
INFO:root:Entering: test
INFO:root:Exiting: test

In [10]: with test_context:
    ...:     pass
    ...: 
INFO:root:Entering: test
INFO:root:Exiting: test

How can I create a class that fails on the second call to the with statement?

Answer 1

I suggest to consider iterable class instead of context manager, like this

class Iterable:
    """Iterable that can be iterated only once."""

    def __init__(self, name):
        self.name = name
        self.it = iter([self])

    def __iter__(self):
        # code to acquire resource
        print('enter')
        yield next(self.it)
        print('exit')
        # code to release resource

    def __repr__(self):
        return f'{self.__class__.__name__}({self.name})'

It can be iterated only one

>>> it = Iterable('iterable')
>>> for item in it:
>>>     print('entered', item)
enter
entered Iterable(iterable)
exit
>>> for item in it:
>>>     print('entered', item)
RuntimeError: generator raised StopIteration

Context manager can be written in this manner:

class Context:
    """Context manager that can be used only once."""

    def __init__(self, name):
        self.name = name
        self.it = iter([self])

    def __enter__(self):
        print('enter')
        return next(self.it)

    def __exit__(self, exc_type, exc, exc_tb):
        print('exit')

    def __repr__(self):
        return f'{self.__class__.__name__}({self.name})'

It works only once

>>> ctx = Context('context')
>>> with ctx as c:
>>>     print('entered', c)
enter
entered Context(context)
exit
>>> with ctx as c:
>>>     print('entered', c)
enter
StopIteration:

Answer 2

Arguably the simplest method is mentioned two paragraphs further down in the documentation you have cited:

Context managers created using contextmanager() are also single use context managers, and will complain about the underlying generator failing to yield if an attempt is made to use them a second time

Here is the corresponding invocation for your example:

>>> from contextlib import contextmanager
>>> @contextmanager
... def track_entry_and_exit(name):
...    print('Entering', name)
...    yield
...    print('Exiting', name)
... 
>>> c = track_entry_and_exit('test')
>>> with c:
...    pass
... 
Entering test
Exiting test
>>> with c:
...    pass
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.9/contextlib.py", line 115, in __enter__
    del self.args, self.kwds, self.func
AttributeError: args

It's even a class although it is written as a function:

>>> type(c)
<class 'contextlib._GeneratorContextManager'>

Answer 3

Here is a possible solution:

from functools import wraps


class MultipleCallToCM(Exception):
    pass


def single_use(cls):
    if not ("__enter__" in vars(cls) and "__exit__" in vars(cls)):
        raise TypeError(f"{cls} is not a Context Manager.")

    org_new = cls.__new__
    @wraps(org_new)
    def new(clss, *args, **kwargs):
        instance = org_new(clss)
        instance._called = False
        return instance
    cls.__new__ = new

    org_enter = cls.__enter__
    @wraps(org_enter)
    def enter(self):
        if self._called:
            raise MultipleCallToCM("You can't call this CM twice!")
        self._called = True
        return org_enter(self)

    cls.__enter__ = enter
    return cls


@single_use
class CM:
    def __enter__(self):
        print("Enter to the CM")

    def __exit__(self, exc_type, exc_value, exc_tb):
        print("Exit from the CM")


with CM():
    print("Inside.")
print("-----------------------------------")
with CM():
    print("Inside.")
print("-----------------------------------")
cm = CM()
with cm:
    print("Inside.")
print("-----------------------------------")
with cm:
    print("Inside.")

output:

Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Traceback (most recent call last):
  File "...", line 51, in <module>
    with cm:
  File "...", line 24, in enter
    raise MultipleCallToCM("You can't call this CM twice!")
__main__.MultipleCallToCM: You can't call this CM twice!

I used a class decorator for it so that you can apply it to other context manager classes. I dispatched the __new__ method and give every instance a flag called __called, then change the original __enter__ to my enter which checks to see if this object has used in a with-statement or not.

How robust is this? I don't know. Seems like it works, I hope it gave an idea at least.