Dump dot file by using modulegraph. It trims unwanted modules in the graph.

dump-module-dep.py

#!/usr/bin/env python3
#
# Author: Taihsiang Ho (tai271828)
# License: BSD 3-clause
#
# This script uses modulegraph to prepare a Graphviz dot file.
#
# The first positional argument is the target python script to be scanned. The
# result of the optional arguments could refer to the --help.
#
# Prequisite python packages:
#   click
#   modulegraph
#
#
# Example (take for the project SOLVCON in my home as an example):
#   ./dump-module-dep.py solvcon/parcel/gas/physics.py --ppath /home/tai271828/work-my-projects/solvcon/solvcon/
#
#   ./dump-module-dep.py solvcon/parcel/gas/physics.py --ppath /home/tai271828/work-my-projects/solvcon/solvcon/ -r ctypes -r numpy -r __future__ -r six
#
#   ./dump-module-dep.py solvcon/parcel/gas/physics.py --ppath /home/tai271828/work-my-projects/solvcon/solvcon/ -r ctypes -r numpy -r __future__ -r six -k solvcon
#
#   ./dump-module-dep.py solvcon/parcel/gas/physics.py --ppath /home/tai271828/work-my-projects/solvcon/solvcon/ -r ctypes -r numpy -r __future__ -r six -k solvcon -w /home/tai271828/work-my-projects/solvcon/solvcon/solvcon/parcel/gas/physics.py
#
#   and then use dot tools to generate a png file
#
#   dot physics.dot -T png -o physics.png && xdg-open physics.png
#
#
import posixpath
import click
from modulegraph.modulegraph import ModuleGraph


def run(script_2run,
        search_path=[], nodes_2remove_prefix=[], nodes_2reserve_prefix=[],
        whitelist=[]):
    # get graph
    mg = ModuleGraph(search_path)
    mg.run_script(script_2run)

    # collect nodes to remove
    # white list (never remove it)
    # keep script of run_script in order not to remove the top of the graph
    # tree
    nodes_2remove = []
    nodes = []
    for node in mg.nodes():
        nodes.append(node.identifier)
        # get top package or module name
        node_id_prefix = node.identifier.split('.')[0]
        if node_id_prefix in nodes_2remove_prefix:
            # positive list to remove
            nodes_2remove.append(node.identifier)
        elif len(nodes_2reserve_prefix) > 0:
            # len > 0 so I won't remove everything
            if node_id_prefix not in nodes_2reserve_prefix:
                # negative list to remove
                nodes_2remove.append(node.identifier)

    # keep nodes; kick out nodes in the nodes_2remove
    for node in whitelist:
        if node in nodes_2remove:
            nodes_2remove.remove(node)

    # do real action to remove nodes
    for node in nodes:
        if node in nodes_2remove:
            mg.removeNode(node)

    # dump the graph, use
    # dot physics.dot -T png -o physics.png
    # to get a png picture
    f = open("physics.dot", "w")
    mg.graphreport(f)
    f.close()


def show_usage():
    click.echo("Please input parameters correctly")
    click.echo("Refer to --help")


@click.command()
@click.argument('pscript', type=str)
@click.option('-p', '--ppath', multiple=True, type=str,
              help='Where the packages to be searched are')
@click.option('-r', '--remove-prefix', multiple=True, type=str,
              help='Node with this prefix will be removed.')
@click.option('-k', '--keep-prefix', multiple=True, type=str,
              help='Node with this prefix will be reserved.')
@click.option('-w', '--whitelist', multiple=True, type=str,
              help='Node with this identifier name will be reserved.')
def gendot(pscript, ppath, remove_prefix, keep_prefix, whitelist):
    if ppath is None or pscript is None:
        show_usage()
        return

    # Always assume 1st ppath is the project root path
    script_2run = posixpath.join(ppath[0], pscript)
    click.echo('Scan script: %s' % script_2run)
    click.echo('Search packages and mouldes in: %s' % ppath)

    # positive list to remove
    nodes_2remove_prefix = remove_prefix
    click.echo('Nodes with prefix to remove: %s' % str(nodes_2remove_prefix))
    # negiaitve list to remove
    nodes_2reserve_prefix = keep_prefix
    click.echo('Nodes with prefix to reserve: %s' % str(nodes_2reserve_prefix))

    # set some default
    if len(ppath) == 1 and len(whitelist) < 1:
        # Please note we always reserve the scanned script itself
        whitelist = [script_2run]
        click.echo('Nodes of whitelist: %s' % str(whitelist))

    run(script_2run, ppath,
        nodes_2remove_prefix, nodes_2reserve_prefix, whitelist)


if __name__ == '__main__':
    gendot()