Commits

Kamil Kisiel committed 3bb747a

Factored out module finding code in to find_module(), cleaned up some imports.

Comments (0)

Files changed (1)

importgraph/__init__.py

 import imp
 import networkx
 import matplotlib.pyplot as plt
+import os
 import sys
 from collections import namedtuple
 
 # TODO: Generate this in some better way. This is only an example
 STDLIB_MODULES = ['os', 'subprocess', 'sys', 'inspect', 'functools', 'collections']
 
+def find_module(module):
+    """Find the file containing a module."""
+    module_filename = None
+    basename = module.replace(".", "/")
+    # Looking for module file; can be either a .py 
+    # file or a python package.
+    for path in sys.path:
+        # TODO: probably better to check if dirname(path)
+        # exists first, then look for specific types of files.
+        module_filename = os.path.join(path, basename + ".py")
+        if os.access(module_filename, os.F_OK):
+            break
+        module_filename = os.path.join(path, basename, "__init__.py")
+        if os.access(module_filename, os.F_OK):
+            break
+        module_filename = os.path.join(path, basename + ".so")
+        if os.access(module_filename, os.F_OK):
+            break
+    return module_filename
+
 class ImportVisitor(ast.NodeVisitor):
 
     def __init__(self, graph, root, ignore=None, visited=None, level=None):
 
         self.visited.append(module)
 
-        basename = module.replace(".", "/")
-        module_filename = None
-
         # TODO add filtering
         display_name = module.split(".", 1)[0]
         self.graph.add_edge(display_name, self.root)
 
-        # Looking for module file; can be either a .py 
-        # file or a python package.
-        import sys, os
-        for path in sys.path:
-            # TODO: probably better to check if dirname(path)
-            # exists first, then look for specific types of files.
-            module_filename = os.path.join(path, basename + ".py")
-            if os.access(module_filename, os.F_OK):
-                break
-            module_filename = os.path.join(path, basename, "__init__.py")
-            if os.access(module_filename, os.F_OK):
-                break
-            module_filename = os.path.join(path, basename + ".so")
-            if os.access(module_filename, os.F_OK):
-                break
+        module_filename = find_module(module)
 
         # Recurse if possible.
         if module not in self.ignore and module_filename: