Is it okay to use the method of a class as contextmanager?

Question

so, here is a simple code:

from contextlib import contextmanager

class A:
    
    def __init__(self): 
        self.a = 1
        return

    @contextmanager
    def a2mode(self):
        self.a = 2
        print(self.a)
        yield
        self.a = 1
        print(self.a)
        return
    
ca = A()
with ca.a2mode() as cm:
    print(ca.a+1)

which prints:

2
3
1

Basically I am using a method of the class to create a context within which the attribute a has a value a=2, instead of its usual value. Based on that some operations are performed that make use of the modified value of a. After operation within that context, the attribute is reset to its original value before the context manager.

Is this ok to do, or can it lead to some unforeseen issues? It works, but I am not sure whether this is "forbidden" code.

Thanks! Best, JZ

Answer 1

That should be fine, though if you want to ensure a gets reset even if the context-managed block raises, you'll want to use try: finally:.

    @contextmanager
    def a2mode(self):
        self.a = 2
        try:
            yield
        finally:
            self.a = 1

Also, if there would be more modes, you'd want to make sure to reset to the previous mode...

    @contextmanager
    def a2mode(self):
        old_a = self.a
        self.a = 2
        try:
            yield
        finally:
            self.a = old_a