More Pythonic Context Manager Wrapper

Question

I wanted to ask this here because what I wanted to originally do felt like a really Pythonic method. I want to be able to use the syntax:

d = {'apple':{'cranberry':{'banana':{'chocolate':[1,2,3,4,5,6]}}},'b':2}
with d['apple']['cranberry']['banana']['chocolate'] as item:
    for i in item:
        print(i)

    item.append('k')

but found that Python doesn't allow for using lists, dicts, etc. as context managers.

So I implemented my own:

def context_wrap(target):
    class ContextWrap:
        def __init__(self, tgt):
            self._tgt = tgt

        def __enter__(self):
            return self._tgt

        def __exit__(self, type, value, traceback):
            pass

    return ContextWrap(target)

if __name__ == '__main__':
    with context_wrap({'a':1}) as item:
        print(item['a'])

    with context_wrap([1,2,3,4,5]) as item:
        print(item[1])

    with context_wrap(3) as item:
        print (item)

In the above code, you can take any random object and wrap it inside an object that acts as a context manager controlling the underlying object. This means that inside any with clause, you can simply use the object with its alias. I feel like it looks a lot cleaner and clearer than something like:

alias = d['apple']['cranberry']['banana']['chocolate']
for i in alias:
    print(i)

alias.append('k')

I wanted to know if there was a more "Pythonic" way to do it. So the improvement that I'm looking for is in terms of better reliance on the standard Python library and/or in terms of my syntax.

Answer 1

I think what you are doing is crazy, but yes, you can use Python’s library functions to clean it up.

from contextlib import contextmanager

@contextmanager
def context_wrap(target):
    yield target

Again, this is busy work.

alias = thing

is clearer, shorter and faster than

with context_wrap(thing) as alias:

Answer 2

AJNeufeld provides the proper way to create small context managers using contextlib.contextmanager. While this is, for sure, the best way to accomplish your goal, I wanted to comment on the general use case of functions like the following:

def context_wrap(target):
    class ContextWrap:
        def __init__(self, tgt):
            self._tgt = tgt

        def __enter__(self):
            return self._tgt

        def __exit__(self, type, value, traceback):
            pass

    return ContextWrap(target)

With this sort of function, you are creating the class, instantiating it, and then throwing away the class definition on each function call. All this needs to be is the class definition, and then use that directly:

class ContextWrap:
    def __init__(self, tgt):
        self._tgt = tgt

    def __enter__(self):
        return self._tgt

    def __exit__(self, type, value, traceback):
        pass


with ContextWrap(target) as cw:
    # do something

Which is much lighter since you aren't re-creating resources all the time.

Any time you find yourself using that pattern, go simpler.