CODEWITHSUNDEEP
python
Jellby
Jellby

Reputation: 2686

Context manager inside a class instance

I would like to create a context manager class that's specific for an instance of another class. I can do it by calling a method that creates a class, but I'm not sure this is the best, nicest way:

class MyClass(object):
    def __init__(self):
        self.level = 0
        self.Nest = self.create_context_manager()
    def inclev(self):
        self.level += 1
    def declev(self):
        self.level -= 1

    def create_context_manager(self):
        self2 = self
        class Nest(object):
            def __init__(self):
                pass
            def __enter__(self):
                self2.inclev()
            def __exit__(self, exc_type, exc_value, traceback):
                self2.declev()
        return Nest

# Manually increase/decrease level
my_instance = MyClass()
print(my_instance.level)
my_instance.inclev()
print(my_instance.level)
my_instance.inclev()
print(my_instance.level)
my_instance.declev()
print(my_instance.level)
my_instance.declev()
print(my_instance.level)

# Use instance-specific context manager
other_instance = MyClass()
print(other_instance.level)
with other_instance.Nest():
    print(other_instance.level)
    with other_instance.Nest():
        print(other_instance.level)
    print(other_instance.level)
print(other_instance.level)

Upvotes: 11

Views: 15120

Answers (1)

ShadowRanger
ShadowRanger

Reputation: 155363

Why do you need a nested class? Just have the main object implement the context management protocol directly:

class MyClass(object):
    def __init__(self):
        self.level = 0
    def inclev(self):
        self.level += 1
    def declev(self):
        self.level -= 1

    __enter__ = inclev  # For __enter__, just alias inclev, no need for wrapper
    def __exit__(self, exc_type, exc_value, traceback):
        self.declev()

Then just use it with:

with other_instance:
    print(other_instance.level)
    with other_instance:
        print(other_instance.level)
    print(other_instance.level)

If you really need the context manager protocol as a constructed thing named Nest, you can still simplify a bit with builtins from contextlib:

from contextlib import contextmanager

class MyClass(object):
    def __init__(self):
        self.level = 0
    def inclev(self):
        self.level += 1
    def declev(self):
        self.level -= 1

    @contextmanager
    def Nest(self):
        self.inclev()
        try:
            yield
        finally:
            self.declev()

Upvotes: 12

Related Questions