diff --git a/pipdeptree.py b/pipdeptree.py index a9672d9..4c36a17 100644 --- a/pipdeptree.py +++ b/pipdeptree.py @@ -106,6 +106,35 @@ def reverse_tree(tree): return rtree +def filter_tree(tree, list_all=True, show_only=None): + + """Filter the tree. + + :param dict tree: the package tree + :param bool list_all: whether to list all the pgks at the root + level or only those that are the + sub-dependencies + :param set show_only: set of select packages to be shown in the + output. This is optional arg, default: None. + :returns: reversed tree + :rtype: dict + + """ + tree = sorted_tree(tree) + branch_keys = set(r.key for r in flatten(tree.values())) + nodes = tree.keys() + node_keys = [node.key for node in nodes] + + if show_only: + node_keys = [node.key for node in nodes + if node.key in show_only or node.project_name in show_only] + elif not list_all: + node_keys = [node.key for node in nodes if node.key not in branch_keys] + + return dict([ + (node, deps) for node, deps in tree.items() if node.key in node_keys]) + + def guess_version(pkg_key, default='?'): """Guess the version of a pkg when pip doesn't provide it @@ -274,35 +303,22 @@ class ReqPackage(Package): 'required_version': self.version_spec} -def render_tree(tree, list_all=True, show_only=None, frozen=False): +def render_tree(tree, frozen=False): """Convert to tree to string representation :param dict tree: the package tree - :param bool list_all: whether to list all the pgks at the root - level or only those that are the - sub-dependencies - :param set show_only: set of select packages to be shown in the - output. This is optional arg, default: None. :param bool frozen: whether or not show the names of the pkgs in the output that's favourable to pip --freeze :returns: string representation of the tree :rtype: str """ - tree = sorted_tree(tree) - branch_keys = set(r.key for r in flatten(tree.values())) nodes = tree.keys() use_bullets = not frozen key_tree = dict((k.key, v) for k, v in tree.items()) get_children = lambda n: key_tree.get(n.key, []) - if show_only: - nodes = [p for p in nodes - if p.key in show_only or p.project_name in show_only] - elif not list_all: - nodes = [p for p in nodes if p.key not in branch_keys] - def aux(node, parent=None, indent=0, chain=None): if chain is None: chain = [node.project_name] @@ -469,8 +485,7 @@ def get_parser(): parser.add_argument('-j', '--json', action='store_true', default=False, help=( 'Display dependency tree as json. This will yield ' - '"raw" output that may be used by external tools. ' - 'This option overrides all other options.' + '"raw" output that may be used by external tools.' )) parser.add_argument('--graph-output', dest='output_format', help=( @@ -491,6 +506,9 @@ def main(): dist_index = build_dist_index(pkgs) tree = construct_tree(dist_index) + show_only = set(args.packages.split(',')) if args.packages else None + tree = filter_tree(tree, list_all=args.all, show_only=show_only) + if args.json: print(jsonify_tree(tree, indent=4)) return 0 @@ -528,10 +546,7 @@ def main(): if args.warn == 'fail' and (conflicting or cyclic): return_code = 1 - show_only = set(args.packages.split(',')) if args.packages else None - tree = render_tree(tree if not args.reverse else reverse_tree(tree), - list_all=args.all, show_only=show_only, frozen=args.freeze) print(tree) return return_code diff --git a/tests/test_pipdeptree.py b/tests/test_pipdeptree.py index 623f361..cdcbeca 100644 --- a/tests/test_pipdeptree.py +++ b/tests/test_pipdeptree.py @@ -2,7 +2,7 @@ import pickle from operator import itemgetter, attrgetter from pipdeptree import (build_dist_index, construct_tree, - DistPackage, ReqPackage, render_tree, + DistPackage, ReqPackage, filter_tree, render_tree, reverse_tree, cyclic_deps, conflicting_deps) @@ -98,7 +98,7 @@ def test_ReqPackage_render_as_branch(): def test_render_tree_only_top(): - tree_str = render_tree(tree, list_all=False) + tree_str = render_tree(filter_tree(tree, list_all=False)) lines = set(tree_str.split('\n')) assert 'Flask-Script==0.6.6' in lines assert ' - SQLAlchemy [required: >=0.7.3, installed: 0.9.1]' in lines @@ -107,7 +107,7 @@ def test_render_tree_only_top(): def test_render_tree_list_all(): - tree_str = render_tree(tree, list_all=True) + tree_str = render_tree(filter_tree(tree, list_all=True)) lines = set(tree_str.split('\n')) assert 'Flask-Script==0.6.6' in lines assert ' - SQLAlchemy [required: >=0.7.3, installed: 0.9.1]' in lines @@ -116,7 +116,7 @@ def test_render_tree_list_all(): def test_render_tree_freeze(): - tree_str = render_tree(tree, list_all=False, frozen=True) + tree_str = render_tree(filter_tree(tree, list_all=False), frozen=True) lines = set() for line in tree_str.split('\n'): # Workaround for https://github.com/pypa/pip/issues/1867 @@ -144,7 +144,7 @@ def test_cyclic_dependencies(): def test_render_tree_cyclic_dependency(): cyclic_pkgs, dist_index, tree = venv_fixture('tests/virtualenvs/cyclicenv.pickle') - tree_str = render_tree(tree, list_all=True) + tree_str = render_tree(tree) lines = set(tree_str.split('\n')) assert 'CircularDependencyA==0.0.0' in lines assert ' - CircularDependencyB [required: Any, installed: 0.0.0]' in lines @@ -154,7 +154,7 @@ def test_render_tree_cyclic_dependency(): def test_render_tree_freeze_cyclic_dependency(): cyclic_pkgs, dist_index, tree = venv_fixture('tests/virtualenvs/cyclicenv.pickle') - tree_str = render_tree(tree, list_all=True, frozen=True) + tree_str = render_tree(tree, frozen=True) lines = set(tree_str.split('\n')) assert 'CircularDependencyA==0.0.0' in lines assert ' CircularDependencyB==0.0.0' in lines