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

Thanks everyone now it works smoothly, I actually wanted with to return something because I wanted to encapsulate the “run” inside the context manager, so I get something as below.

Moreover, now I store the old sys.argv and restore it on exit, probably not fundamental but still a nice thing to do I think..

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.old_argv = argv[:]
        self.output_file = output_file
        self.progname = self.old_argv[1]

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

    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)
        sys.argv = self.old_argv[:]

    def run(self):
        code = compile(open(self.progname).read(), self.progname, 'exec')
        exec(code)


def main_context():
    cl = CollectImports()

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