Writing a Python class that can only be used as a context manager [duplicate]

问题

Is there a way in Python to write a class that will error unless it's used with a with statement?

# Okay:
with Foo() as f1:
    f1.func1()
    f1.func2()

# Not okay:
f2 = Foo()
f2.func1()

I can do it manually: have __enter__ set a flag and have every other method check for that flag. But is there a nicer way to do it?

Here's code for the not-so-since way:

class Foo(object):
    def __init__(self):
        self._entered = False

    def __enter__(self):
        self._entered = True
        return self

    def _verify_entered(self):
        if not self._entered:
            raise Exception("Didn't get call to __enter__")

    def __exit__(self, typ, val, traceback):
        self._verify_entered()
        print("In __exit__")

    def func1(self):
        self._verify_entered()
        # do stuff ...

    def func2(self):
        self._verify_entered()
        # do other stuff

回答1:

Technically, I think that agf has it right in the sense that you could use a metaclass to automate the stuff. However, if I understand the fundamental motivation behind it correctly, I'd suggest a different way.

Suppose you have a Payload class you want to protect via a context manager. In this case, you simply create a context manager that returns it:

# This should go in a private module.
class Payload(object):
    def __init__(self):
        print 'payload ctor'

# This should go in the public interface.
class Context(object):
    def __init__(self):
        # Set up here the parameters.
        pass

    def __enter__(self):
        # Build & return a Payload object
        return Payload()

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Cleanup
        pass

with Context() as f:
    # f here is a Payload object.

If you hide Payload in a private module, you're good.

回答2:

You can have your __enter__ method return a different object than self if you don't want a user to be able to call methods on the context manager object itself.

class Foo(object):
    def __enter__(self):
        print("In __enter__")
        return Bar()

    def __exit__(self, typ, val, traceback):
        print("In __exit__")

class Bar(object):
    def func1(self):
        print("In func1")

    def func2(self):
        print("In func2")

You can of course have the Foo and Bar classes more tied together than I have for this example. For instance, the Foo class could pass itself to the Bar constructor in __enter__.

Calling Foo().func1() will not work for the obvious reason that Foo doesn't have any such method. If you wanted Bar not to be visible to the user, you could prefix its name with an underscore (hinting that it's internal) or even nest it within the Foo class (or even the __enter__ method, if you really want to be extreme).

来源：https://stackoverflow.com/questions/30673402/writing-a-python-class-that-can-only-be-used-as-a-context-manager

标签

python

contextmanager