Is a context manager right for this job?
-
07-03-2021 - |
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()
Solution 3
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()
OTHER TIPS
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 ofwith
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)
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)