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()
snakefood
, which also does something similar? \$\endgroup\$ – Graipher Aug 19 '18 at 18:04