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

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

CodePudding user response：

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

CodePudding user response：

It might be simpler to just define __enter__ and __exit__ explicitly, as __exit__ receives any exception raised in the with statement as an argument.

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

    def __enter__(self):
        self.a = 2
        print(self.a)

    def __exit__(self, *args):
        self.a = 1
        print(self.a)
        # I don't know the most idiomatic way of checking this...
        if args != (None, None, None):
            exc_type, exc_value, traceback = args
            ...
            # Return a true value to prevent the exception from propagating