Don't setdefault in mymf.import_module if import_hook hasn't been run.

[depgraph.git] / py2depgraph.py

diff --git a/py2depgraph.py b/py2depgraph.py
old mode 100644 (file)
new mode 100755 (executable)
index 5e97e06..fa1b930
--- a/py2depgraph.py
+++ b/py2depgraph.py
@@ -1,4 +1,7 @@
-# Copyright 2004 Toby Dickenson
+#!/usr/bin/env python
+#
+# Copyright 2004      Toby Dickenson
+# Copyright 2008-2011 W. Trevor King
 #
 # Permission is hereby granted, free of charge, to any person obtaining
 # a copy of this software and associated documentation files (the
@@ -19,44 +22,72 @@
 # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import sys, pprint
+"""Extract a tree of module imports from a Python script.
+"""
+
+import logging
 import modulefinder
 
+LOG = logging.getLogger('py2depgraph')
+LOG.setLevel(logging.DEBUG)
+_STREAM_HANDLER = logging.StreamHandler()
+_STREAM_HANDLER.setLevel(logging.CRITICAL)
+_STREAM_HANDLER.setFormatter(
+    logging.Formatter('%(levelname)s - %(message)s'))
+LOG.addHandler(_STREAM_HANDLER)
+
+
 class mymf(modulefinder.ModuleFinder):
     def __init__(self,*args,**kwargs):
         self._depgraph = {}
         self._types = {}
+        self._paths = {}
         self._last_caller = None
         modulefinder.ModuleFinder.__init__(self,*args,**kwargs)
-        
-    def import_hook(self, name, caller=None, fromlist=None):
+
+    def import_hook(self, name, caller=None, fromlist=None, level=-1):
         old_last_caller = self._last_caller
         try:
             self._last_caller = caller
-            return modulefinder.ModuleFinder.import_hook(self,name,caller,fromlist)
+            return modulefinder.ModuleFinder.import_hook(self,name,caller,fromlist, level)
         finally:
             self._last_caller = old_last_caller
-            
+        LOG.debug('after import_hook(%s, %s, %s, %s), last_caller is %s'
+                  % (name, caller, fromlist, level, self._last_caller))
+
     def import_module(self,partnam,fqname,parent):
         r = modulefinder.ModuleFinder.import_module(self,partnam,fqname,parent)
-        if r is not None:
+        if None not in [r, self._last_caller]:
             self._depgraph.setdefault(self._last_caller.__name__,{})[r.__name__] = 1
         return r
-    
+
     def load_module(self, fqname, fp, pathname, (suffix, mode, type)):
         r = modulefinder.ModuleFinder.load_module(self, fqname, fp, pathname, (suffix, mode, type))
         if r is not None:
             self._types[r.__name__] = type
+            self._paths[r.__name__] = pathname
         return r
-        
-        
-def main(argv):    
-    path = sys.path[:]
-    debug = 0
-    exclude = []
-    mf = mymf(path,debug,exclude)
-    mf.run_script(argv[0])
-    pprint.pprint({'depgraph':mf._depgraph,'types':mf._types})
-    
+
+
 if __name__=='__main__':
-    main(sys.argv[1:])
\ No newline at end of file
+    from optparse import OptionParser
+    from pprint import pprint
+    import sys
+
+    usage = '%prog [options] path/to/script.py'
+    p = OptionParser(usage=usage, description=__doc__)
+    p.add_option('-v', '--verbose', default=0, action='count',
+                 help='Increment verbosity.')
+
+    options,args = p.parse_args()
+
+    log_level = logging.CRITICAL - 10*options.verbose
+    _STREAM_HANDLER.setLevel(log_level)
+
+    script = args[0]
+
+    debug = options.verbose
+    debug = 0
+    mf = mymf(path=sys.path[:], debug=debug, excludes=[])
+    mf.run_script(script)
+    pprint({'depgraph':mf._depgraph,'types':mf._types,'paths':mf._paths})