Making a graph of the import structure of a program

Question

When I read other people's code, it's sometimes difficult to understand the import structure. I thought making a graph of the import structure of a given python program would be a fun programming task, and maybe would result in a cool visualization. Currently the program included only works for modules with a .__file_ attribute. After the code is reviewed I'll add more functionality. I know some methods in the class below could be written recursively, but they would quickly reach the max recursion depth. I'm looking for comments on style, logic, and structure. Also, instead of creating a file on your filesystem (if you decide to run the code), I chose to let you create a custom file with a few imports named 'your_test_file.py'. I imported os and argparse in my test file, and the program worked fine.

import matplotlib.pyplot as plt
from collections import defaultdict
import networkx as nx
import sys

class ImportGraph(object):

    def __init__(self, filename):
        self.base_filename = filename
        relationship_dict = defaultdict(list) #key: base_filename. Value(s): files that base_filename imports
        module_dict = defaultdict(list) # module and path_to_module
        rd, md = self._parse_file(self.base_filename, relationship_dict, module_dict)

        for i in range(30): 
            rd, md = self._scan_dict(rd, md)

        self.final = self._collapse_dictionaries(rd, md)


    def _collapse_dictionaries(self, relationship_dict, module_dict):
        ''' Changes the keys of relationship dict from full paths to module names
        '''
        new_dict = defaultdict(list)
        new_dict[self.base_filename] = relationship_dict[self.base_filename]
        for key in relationship_dict:
            for module_name in module_dict:
                if module_name in key:
                    new_dict[module_name] = relationship_dict[key]

        return new_dict

    def _scan_dict(self, relationship_dict, module_dict):
        try:
            for filename in relationship_dict.copy():
                imported_modules = relationship_dict[filename]
                for module in imported_modules:
                    try:
                        rd, md = self._parse_file(module_dict[module], relationship_dict, module_dict)
                    except UnicodeDecodeError as e:
                        continue
            return rd, md
        except UnboundLocalError:
            print("Module's path has a .so suffix, quitting program. Try again with a different module!")
            print("Module path:", module_dict[module])
            sys.exit()


    def _parse_file(self, filename, relationship_dict, module_dict):
        ''' 
        Parses file for imports and returns a relationship dict containing the import structure
        as well as a module dict containing the path to each module the file imports.
        Need to go through every line in each file, saving the module_name, module mnemonic, 
        and corresponding path
        '''
        in_comment = False
        if type(filename) == tuple:
            print(filename)
            return relationship_dict, module_dict
        with open(filename, 'r') as f:
            lines = f.readlines()
            for line in lines:
                if '"""' in line and not in_comment:
                    in_comment = True
                    continue
                if "'''" in line and not in_comment:
                    in_comment = True
                    continue
                if '"""' in line and in_comment:
                    in_comment = False

                if "'''" in line and in_comment:
                    in_comment = False

                if 'import' in line and '#' not in line and not in_comment:         
                    try:
                        module_names = self._parse_line_for_module_name(line)
                        module_mnemonics = self._parse_line_for_module_mnemonic(line)
                        if module_mnemonics is not None:
                            for i, module_mnemonic in enumerate(module_mnemonics):
                                path_dir = {}
                                line = line.lstrip().replace('\n', '')
                                exec(line + '\npath = {}.__file__'.format(module_mnemonic), globals(), path_dir)
                                path_to_module = path_dir['path']                                   
                                if module_names[i] not in relationship_dict[filename]: 
                                    module_dict[module_names[i]] = path_to_module 
                                    #print("Module {} imported by file: {}".format(module_names[i], filename))
                                    relationship_dict[filename].append(module_names[i])


                    except Exception as e:
                        continue
        return relationship_dict, module_dict


    def _get_path_homemade_module(self):
        ''' will take care of homemade modules w/o a .__file__ attribute '''


    def _parse_line_for_module_name(self, line):
        '''
        This should return the 'official' name of the module:
        instead of returning plt for line "import matplotlib.pyplot as plt", return matplotlib.pyplot
        matplotlib.pyplot will become a node in the directed graph
        '''
        if ',' in line and 'from' not in line:
            '''ex: import socket, math, struct, time, os, fnmatch, array, sys, errno'''
            sub_line = line[line.find('t')+1:].replace(' ', '').replace("\n", '')
            return sub_line.split(',')
        if 'from' in line:
            sub_line = line[line.find("from")+5:line.find("import")]
            if '.' in sub_line:
                sub_line = sub_line[:sub_line.find('.')]
        elif 'as' in line:
            sub_line = line[7:line.find('as')-1].replace(" ", '').replace('\n', '')
        else:
            sub_line = line[line.find('t')+1:].replace(" ", '').replace('\n', '')
        return [sub_line]

    def _parse_line_for_module_mnemonic(self, line):
        ''' 3 cases: 
            from module import submodule1, submodule2, ...
            import module
            import module as module
        '''
        if ',' in line and 'from' not in line:
            '''ex: import socket, math, struct, time, os, fnmatch, array, sys, errno'''
            sub_line = line[line.find('t')+2:].replace(' ', '').replace("\n", '')
            return sub_line.split(',')

        if 'importlib' in line:
            return 'importlib'

        if 'from' in line:
            sub_line = line[line.find('from')+4:line.find("import")-1].replace(" ", '')

            if sub_line.startswith('.'): #this throws an error on exec() call
                return None

        elif 'as' in line:
            sub_line = line[line.find('as')+2:].replace(" ", '').replace('\n', '')
        else:
            sub_line = line[line.find('t')+1:].replace(" ", '').replace('\n', '')

        return [sub_line]


def main(filename):
    it = ImportTree(filename)

if __name__ == '__main__':

    base_filename = 'your_test_file.py'
    ig = ImportGraph(base_filename)
    colors = []
    final_relationship_dict = ig.final
    G = nx.Graph()

    for key in final_relationship_dict:
        if key == base_filename:
            colors.append('blue')
        else:
            colors.append('red')

        G.add_node(key)

        for entry in final_relationship_dict[key]:
            G.add_edge(key, entry)


    nx.draw(G, node_color = colors, with_labels=True)
    plt.show()

Answer 1

1. If a Python module contains no imports, then the code in the post fails with an error like this:

   Module's path has a .so suffix, quitting program. Try again with a different module!
   Traceback (most recent call last):
     File "cr201985.py", line 41, in _scan_dict
       return rd, md
   UnboundLocalError: local variable 'rd' referenced before assignment
   
   During handling of the above exception, another exception occurred:
   
   Traceback (most recent call last):
     File "<stdin>", line 1, in <module>
     File "cr201985.py", line 15, in __init__
       rd, md = self._scan_dict(rd, md)
     File "cr201985.py", line 44, in _scan_dict
       print("Module path:", module_dict[module])
   UnboundLocalError: local variable 'module' referenced before assignment
   

   Notice that the error message (about the .so suffix) is incorrect.

2. The methods _parse_line_for_module_name and _parse_line_for_module_mnemonic don't use the self argument, so they don't need to be methods, they could just be plain functions. Plain functions are easier to develop and test than methods, because you can just call them with your test data, you don't have to create an object first.

3. Here's how you use the ImportGraph class:

   ig = ImportGraph(base_filename)
   final_relationship_dict = ig.final
   

   After this, ig is not needed any more (and it doesn't have any other public methods or attributes other than final, so you couldn't use it for anything else even if you wanted to). This suggests that you don't actually need a class here, and what you actually want is a function that takes the base filename and returns the module relationship dictionary.

   So instead of a class, use a function; instead of attributes, use local variables; and instead of private methods, use locally defined functions. Like this:

   def import_graph(base_filename):
       "Return import relationship mapping starting at base_filename."
       relationship_dict = defaultdict(list)
       module_dict = defaultdict(list)
   
       def parse_file(filename):
           # ... update relationship_dict and module_dict here ...
   
       def scan_dict():
           # ... update relationship_dict and module_dict here ...
   
       # ... call parse_file and scan_dict here ...
   
       return relationship_dict
   

   See Jack Diederich's talk "Stop Writing Classes" for more about this issue.

4. The strategy for finding all the imports is to call the _scan_dict method thirty times in a loop:

   for i in range(30): 
       rd, md = self._scan_dict(rd, md)
   

   Each time round this loop, _scan_dict looks at all the modules that have been discovered so far, and parses the corresponding files to look for more import statements. The trouble with this approach is that each imported file gets parsed up to thirty times!

   The way to avoid duplicated work here is to maintain a collection of files that have been discovered but not yet parsed:

   unparsed_files = {base_filename}
   

   and a collection of files that have been parsed:

   parsed_files = set()
   

   and then the main loop can work something like this:

   while unparsed_files:
       filename = unparsed_files.pop()
       parsed_files.add(filename)
       parse_file(filename)
   

   and when an import statement is discovered, you can check whether the imported file has already been parsed:

   if filename not in parsed_files:
       unparsed_files.add(filename)
   

   This approach means that each file gets parsed exactly once, and there is no magic number "30", so import chains of arbitrary depth can be processed.

   (This is similar to the breadth-first search algorithm, except that in this case you don't care about the order in which you explore the graph, so unparsed_files can be a set instead of a queue.)

5. The purpose of _parse_file is to parse Python source code and identify import statements. It works by iterating over the lines looking for the string import, together with some heuristics to try to discard strings and comments. Unfortunately the heuristics don't cover all the cases, and it is easy to fool them. Here are four imports which are all ignored by _parse_file:

   import heapq # comment on the same line as import
   
   # ''' triple quotes are commented out
   import heapq
   # '''
   
   '''triple quotes finish on same line they started'''
   import heapq
   ''''''
   
   # line continuation
   import \
       heapq
   

   There is really no substitute for actually parsing the file. Luckily, Python has a built-in parser in the ast module, so getting it right is easy. The way you do it is to call ast.parse to get an abstract syntax tree:

   import ast
   with open(filename, 'rb') as f:
       tree = ast.parse(f.read(), filename)
   

   (Notice that you must open the file in binary mode — this allows the parser to handle file encoding declarations, and so avoid the UnicodeDecodeError that you had trouble with.)

   Then you can walk the tree looking for import ... and from ... import ... statements. This is most easily done by subclassing the ast.NodeVisitor class, for example:

   class ImportVisitor(ast.NodeVisitor):
       "AST visitor that prints the import statements."
       def visit_Import(self, node):
           print('import', ', '.join(alias.name for alias in node.names))
   
       def visit_ImportFrom(self, node):
           print('from', node.module, 'import',
                 ', '.join(alias.name for alias in node.names))
   

   and then visit the tree that you parsed earlier:

   visitor = ImportVisitor()
   visitor.visit(tree)
   

   You'll see that this approach is much simpler than _parse_line_for_module_name and _parse_line_for_module_mnemonic, and it copes with all the difficult cases that I noted above.

   (Of course, in a real program you wouldn't print the imports, you'd accumulate them for further processing, but printing is a simple way to demonstrate the approach.)