Is a context manager right for this job?

Question

The code pasted below does the following:

• creates an import hook
• creates a context manager which sets the meta_path and cleans on exit.
• dumps all the imports done by a program passed in input in imports.log

Now I was wondering if using a context manager is a good idea in this case, because actually I don't have the standard try/finally flow, but just a set up and clean up.

Another thing — with this line:

with CollectorContext(cl, sys.argv, 'imports.log') as cc:

does cc become None? Shouldn't it be a CollectorContext object?

from __future__ import with_statement
import os
import sys

class CollectImports(object):
    """
    Import hook, adds each import request to the loaded set and dumps
    them to file
    """

    def __init__(self):
        self.loaded = set()

    def __str__(self):
        return str(self.loaded)

    def dump_to_file(self, fname):
        """Dump the loaded set to file
        """
        dumped_str = '\n'.join(x for x in self.loaded)
        open(fname, 'w').write(dumped_str)

    def find_module(self, module_name, package=None):
        self.loaded.add(module_name)


class CollectorContext(object):
    """Sets the meta_path hook with the passed import hook when
    entering and clean up when exiting
    """

    def __init__(self, collector, argv, output_file):
        self.collector = collector
        self.argv = argv
        self.output_file = output_file

    def __enter__(self):
        self.argv = self.argv[1:]
        sys.meta_path.append(self.collector)

    def __exit__(self, type, value, traceback):
        # TODO: should assert that the variables are None, otherwise
        # we are quitting with some exceptions
        self.collector.dump_to_file(self.output_file)
        sys.meta_path.remove(self.collector)


def main_context():
    cl = CollectImports()

    with CollectorContext(cl, sys.argv, 'imports.log') as cc:
        progname = sys.argv[0]
        code = compile(open(progname).read(), progname, 'exec')
        exec(code)


if __name__ == '__main__':
    sys.argv = sys.argv[1:]
    main_context()

Answer 1

I think this concept is ok. As well, I don't see any reasons against having the clean-up stuff in a finally: clause, so the context manager fits perfectly.

Your cc is None, because you told it to be so.

If you don't want that, change your __enter__ method to return something else:

The value returned by this method is bound to the identifier in the as clause of with statements using this context manager.

def __enter__(self):
    self.argv = self.argv[1:]
    sys.meta_path.append(self.collector)
    return self
    # or
    return self.collector
    # or
    return "I don't know what to return here"

and then

with CollectorContext(cl, sys.argv, 'imports.log') as cc:
    print cc, repr(cc) # there you see what happens.
    progname = sys.argv[0]
    code = compile(open(progname).read(), progname, 'exec')
    exec(code)

Answer 2

If you always want the cleanup to occur, you should use a context manager. I'm not sure where you use try..finally if you implement the context manager using the low-level special methods. If you use the @contextmanager decorator, you code the context manager in a "natural" way, so that's where you use try..finally instead of getting the exception as a parameter.

Also, cc will be the value you return from __enter__(). In your case, None. The way I understand the context manager design is that the return value is the "context". What the context manager does is set up and clean up contexts in which something else happens. E.g. a database connection will create transactions, and database operations happen in the scope of those transactions.

That said, the above is just there to provide maximum flexibility. There's nothing wrong with just creating a context (that manages itself) directly and returning self, or even not returning anything if you don't need to use the context value inside the with. Since you don't use cc anywhere, you could just do and not worry about the return value:

with CollectorContext(cl, sys.argv, 'imports.log'):
        progname = sys.argv[0]
        code = compile(open(progname).read(), progname, 'exec')
        exec(code)