1. Kamil Kisiel
  2. importgraph

Commits

Kamil Kisiel  committed 7410231

Added --ignore-stdlib and --ignore options.

  • Participants
  • Parent commits 6785956
  • Branches default

Comments (0)

Files changed (1)

File importgraph.py

View file
 ModuleInfo = namedtuple("ModuleInfo", "name,path")
 MAX_DEPTH = 5
 
+# TODO: Generate this in some better way. This is only an example
+STDLIB_MODULES = ['os', 'subprocess', 'sys', 'inspect', 'functools', 'collections']
+
 class ImportVisitor(ast.NodeVisitor):
 
-    def __init__(self, graph, root, visited=None, level=None):
+    def __init__(self, graph, root, ignore=None, visited=None, level=None):
         self.graph = graph
         self.root = root
+        self.ignore = ignore or set()
         self.visited = visited or []
         self.level = level or 1
 
                 break
 
         # Recurse if possible.
-        if module_filename:
+        if module not in self.ignore and module_filename:
             if module_filename.endswith(".py"):
-                print "processing", module_filename
+                print "processing", module, module_filename
                 try:
                     with open(module_filename) as code:
                         tree = ast.parse(code.read())
-                    visitor = ImportVisitor(self.graph, display_name, self.visited, self.level + 1)
+                    visitor = ImportVisitor(self.graph, 
+                                            display_name, 
+                                            visited=self.visited, 
+                                            ignore=self.ignore,
+                                            level=(self.level + 1))
                     visitor.visit(tree)
                 except SyntaxError, e:
                     pass
 def main():
     argparser = argparse.ArgumentParser()
     argparser.add_argument("pyfile", help="Source file to analyze")
+    argparser.add_argument("--ignore-stdlib", action='store_true',
+                           help="Don't inspect modules in Python's standard library")
+    argparser.add_argument("--ignore", nargs="+", default=[], metavar="MODULE",
+                           help="List of modules names to ignore, spearated by a space")
     args = argparser.parse_args(sys.argv[1:])
+
+    # Load the ast from the file being analyzed
     with open(args.pyfile) as code:
         tree = ast.parse(code.read())
+
+    # Create the intial graph
     graph = networkx.DiGraph()
     graph.add_node(args.pyfile)
-    visitor = ImportVisitor(graph, args.pyfile)
+
+    # Anaylze the AST
+    ignore = set()
+    if args.ignore_stdlib:
+        ignore.update(STDLIB_MODULES)
+    ignore.update(args.ignore)
+
+    visitor = ImportVisitor(graph, args.pyfile, ignore=ignore)
     visitor.visit(tree)
+
+    # Plot output the results
     networkx.draw_graphviz(graph)#, prog='dot')
     plt.show()