conda-tree.py

#!/usr/bin/env python
from __future__ import print_function
import argparse
import json
import os
import sys
import subprocess
import colorama
from colorama import Style
colorama.init()

import conda.exports
import conda.api
import conda.base.context
import networkx

__version__ = '1.1.1'

# The number of spaces
TABSIZE = 3

def get_local_cache(prefix):
    return conda.exports.linked_data(prefix=prefix)

def get_package_key(cache, package_name):
    ks = list(filter(lambda i: cache[i]['name'] == package_name, cache))
    return ks[0]

def make_cache_graph(cache):
    g = networkx.DiGraph()
    for k in cache.keys():
        n = cache[k]['name']
        v = cache[k]['version']
        g.add_node(n, version=v)
        for j in cache[k]['depends']:
            n2 = j.split(' ')[0]
            v2 = j.split(' ')[1:]
            g.add_edge(n, n2, version=v2)
    return(g)

def print_graph_dot(g, exclude_pkgs=set()):
    print("digraph {")
    for k,v in g.edges():
        if k not in exclude_pkgs and v not in exclude_pkgs:
            print(f"  \"{k}\" -> \"{v}\"")
    print("}")

def remove_from_graph(g, node, _cache=None):
    if _cache is None: _cache = {}
    if node not in _cache:
        _cache[node] = True
        for k,v in g.out_edges(node):
            g = remove_from_graph(g, v, _cache)
    if node in g: g.remove_node(node)
    return(g)

def print_dep_tree(g, pkg, prev, state):
    # Unpack the state data
    down_search, args = state["down_search"], state["args"]
    indent = state["indent"]
    empty_cols, is_last = state["empty_cols"], state["is_last"]

    s = ""                          # String to print
    v = g.nodes[pkg].get('version') # Version of package

    # Create list of edges
    edges = g.out_edges(pkg) if down_search else g.in_edges(pkg)
    e = [i[1] for i in edges] if down_search else [i[0] for i in edges]
    # Maybe?: Sort edges in alphabetical order
    # e = sorted(e, key=(lambda i: i[1] if down_search else i[0]))

    if len(args.exclude) > 0:
        for p in args.exclude:
            state['tree_exists'].add(p)

    dependencies_to_hide = (True # We hide dependencies if...
        if ((pkg in state["tree_exists"] and not args.full)
            # Package already displayed and '--full' not used.
            or (args.full and pkg in state["tree_exists"] and pkg in state['pkgs_with_cycles']))
            # or, if '--full' is used but the package is part of a cyclic sub-graph
        else False)
    will_create_subtree = (True if len(e) >= 1 else False)
    if len(e) > 0: state["tree_exists"].add(pkg)

    # If the package is a leaf
    if indent == 0:
        if v is not None:
            s += f"{pkg}=={v}\n"
        else:
            s += pkg
    # Let's print the branch
    else:
        # Finding requirements for package from the parent
        # (or child, if we are running the 'whoneeds' subcommand)
        requirement = (', '.join(g.edges[prev, pkg]['version'])
            if down_search else ', '.join(g.edges[pkg, prev]['version']))
        r = 'any' if requirement == '' else requirement
        # Preparing the prepend string
        br = ('└─' if is_last else '├─')
        # Optional: + ('┬' if will_create_subtree else '─')
        i = ""
        for x in range(indent):
            if x == 0:
                i += ' ' * 2
            elif x in empty_cols:
                i += ' ' * TABSIZE
            else:
                i += ('│' + (' ' * (TABSIZE - 1)))
        if v is not None:
            s += f"{i}{br} {pkg}{Style.DIM} {v} [required: {r}]{Style.RESET_ALL}\n"
        else:
            s += f"{i}{br} {pkg}{Style.DIM} [required: {r}]{Style.RESET_ALL}\n"
        if dependencies_to_hide:
            state["hidden_dependencies"] = True
            will_create_subtree = False
            # We do not print these lines if:
            # python and conda dependencies if '-small' on
            if pkg in args.exclude:
                pass
            else:
                br2 = ' ' if is_last else '│'
                word = "dependencies" if down_search else "dependent packages"
                s += f"{i}{br2}  {Style.DIM}└─ {word} of {pkg} displayed above{Style.RESET_ALL}\n"
        else:
            if len(e) > 0: state["tree_exists"].add(pkg)

    # Print the children
    if will_create_subtree:
        state["indent"] += 1
        for pack in e:
            if state["is_last"]: state["empty_cols"].append(indent)
            state["is_last"] = False if e[-1] != pack else True
            tree_str, state = print_dep_tree(g, pack, pkg, state)
            s += tree_str
    # If this is the last of its subtree to be printed
    if is_last and indent != 0:
        state["indent"] -= 1
        if indent in empty_cols: state["empty_cols"].remove(indent)
        state["is_last"] = False
    return s, state

def get_pkg_files(prefix):
    pkg_files = set()
    for p in conda.api.PrefixData(prefix).iter_records():
        for f in p['files']:
            pkg_files.add(f)
    return pkg_files

# check if dir is internal of conda
def is_internal_dir(prefix,path):
    for t in ['pkgs','conda-bld','conda-meta','locks','envs']:
        if path.startswith(os.path.join(prefix,t)): return True
    return False

def find_who_owns_file(prefix, target_path):
    for p in conda.api.PrefixData(prefix).iter_records():
        for f in p['files']:
            if target_path in f or f in target_path:
                print(f'{p["name"]}\t{f}')

def find_unowned_files(prefix):
    pkg_files = get_pkg_files(prefix)

    for root, dirs, files in os.walk(prefix):
        if is_internal_dir(prefix,root):
            continue

        for f in files:
            f0 = os.path.join(root,f)
            f1 = f0.replace(prefix, "", 1).lstrip(os.sep)
            if f1 not in pkg_files:
                print(f0)

def is_node_reachable(graph, source, target):
    if isinstance(source, list):
        for s in source:
            if is_node_reachable(graph, s, target):
                return True
        return False
    else:
        try:
            paths = networkx.shortest_path(graph, source, target)
            return len(paths)>0
        except:
            return False

def print_pkgs(pkgs, with_json=False):
    if with_json:
        print(json.dumps(pkgs))
    else:
        for p in pkgs:
            print(p)

def find_reachable_pkgs(graph, pkg, down_search=True, exclude_pkgs=set()):
    if down_search:
        paths = networkx.shortest_path(graph, source=pkg)
    else:
        paths = networkx.shortest_path(graph, target=pkg)

    reachable_pkgs = []
    for k, v in paths.items():
        if len(exclude_pkgs.intersection(v)) > 0 and k not in exclude_pkgs:
            pass # remove paths with excluded packages
        elif k != pkg:
            reachable_pkgs.append(k)

    return reachable_pkgs

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p','--prefix',
        help='full path to environment location (i.e. prefix)',
        default=None)
    parser.add_argument('-n','--name',
        help='name of environment',
        default=None)
    parser.add_argument('-V','--version',
        action='version',
        version='%(prog)s '+__version__)

    subparser = parser.add_subparsers(dest='subcmd')

    format_args = argparse.ArgumentParser(add_help=False)
    # Arguments for "rec_or_tree" commands
    # Subcommands that can yield direct dependencies, recursive dependencies, or a tree view
    rec_or_tree = format_args.add_mutually_exclusive_group(required=False)
    rec_or_tree.add_argument('-t', '--tree',
        help='show dependencies of dependencies in tree form',
        action="store_true", default=False)
    rec_or_tree.add_argument('--dot',
         help='print a graphviz dot graph notation',
         action="store_true", default=False)
    rec_or_tree.add_argument('--json',
        help='print packages in json format',
        default=False, action="store_true")

    # Arguments for "package_cmds" commands
    # Subcommands that deal with the dependencies of packages
    package_cmds = argparse.ArgumentParser(add_help=False, parents=[format_args])
    package_cmds.add_argument('package', help='the target package')
    package_cmds.add_argument('-r','--recursive',
        help='show dependencies of dependencies',
        action='store_true',
        default=False)

    # Arguments for "hiding" commands
    # Subcommands that enable users to hide a part of the result
    hiding_cmds = argparse.ArgumentParser(add_help=False)
    hiding_cmds.add_argument(
        '--exclude',
        help='comma separated list of packages to exclude dependencies from tree, ' +
        'can be specified multiple times',
        default=[],
        action='append'
    )
    hiding_cmds.add_argument('--small',
        help="don't include dependencies for conda and python. alias for --exclude conda,python",
        default=False, action='store_true')
    hiding_cmds.add_argument('--full',
        help='shows the complete dependency tree, ' +
              'with all the redundancies that it entails',
        default=False, action='store_true')

    # Definining the simple subcommands
    lv_cmd = subparser.add_parser('leaves',
         help='shows leaf packages')
    lv_cmd.add_argument('--export',
        help='export leaves dependencies',
        default=False, action='store_true')
    lv_cmd.add_argument('--with-cycles',
        help='include orphan cycles',
        default=False, action='store_true')
    lv_cmd.add_argument('--json',
        help='print packages in json format',
        default=False, action="store_true")

    subparser.add_parser('cycles', help='shows dependency cycles')

    # Defining the complex subcommands
    subparser.add_parser('whoneeds',
        help='shows packages that depends on this package',
        parents=[package_cmds, hiding_cmds])
    subparser.add_parser('depends',
        help='shows this package dependencies',
        parents=[package_cmds, hiding_cmds])
    subparser.add_parser('deptree',
        help="shows the complete dependency tree",
        parents=[format_args, hiding_cmds])
    subparser.add_parser('unowned-files',
        help='shows files that are not owned by any package')
    subparser.add_parser('who-owns',
        help='find which package owns a given file').add_argument('file',help='a file path or substring of the target file')

    args = parser.parse_args()

    # Allow user to specify name, but check the environment for an
    # existing CONDA_EXE command.  This allows a different conda
    # package to be installed (and imported above) but will
    # resolve the name using their expected conda.  (The imported
    # conda here will find the environments, but might not name
    # them as the user expects.)
    if args.prefix is None:
        _conda = os.environ.get('CONDA_EXE', 'conda')
        _info = json.loads(subprocess.check_output(
            [_conda, 'info', '-e', '--json']))
        if args.name is None:
            if _info['active_prefix'] is not None:
                args.prefix = _info['active_prefix']
            else:
                args.prefix = _info['default_prefix']
        else:
            args.prefix = conda.base.context.locate_prefix_by_name(
                name=args.name, envs_dirs=_info['envs_dirs'])

    l = get_local_cache(args.prefix)
    g = make_cache_graph(l)

    ######
    # Helper functions for subcommands
    ######
    def get_leaves(graph):
        return list(map(lambda i:i[0],(filter(lambda i:i[1]==0,graph.in_degree()))))

    def get_leaves_plus_cycles(graph):
        lv = get_leaves(graph)
        for pks in networkx.simple_cycles(g):
            if is_node_reachable(g, lv, pks[0]):
                    pass
            else:
                 lv.append(pks[0])
        return lv

    def get_cycles(graph):
        s = ""
        for i in networkx.simple_cycles(graph):
            s += " -> ".join(i)+" -> "+i[0] + "\n"
        return s

    def pkgs_with_cycles(graph):
        return set(sum(networkx.simple_cycles(graph), []))

    # Default state for the recursive tree function
    state = {'down_search': True, 'args': args, 'indent': 0, 'indent': 0,
             'empty_cols': [], 'is_last': False, 'tree_exists': set(),
             'hidden_dependencies': False, 'pkgs_with_cycles': pkgs_with_cycles(g)}

    if args.subcmd in ['depends','whoneeds','deptree']:
        if len(args.exclude) > 0:
            ex = []
            for i in args.exclude:
                for j in i.split(','):
                    ex.append(j)
            args.exclude = ex
        if args.small:
            args.exclude.extend(['conda','python'])

    if args.subcmd == 'cycles':
        print(get_cycles(g), end='')

    elif args.subcmd in ['depends', 'whoneeds']:
        # This variable defines whether we are searching down the dependency
        # tree, or if rather we are looking for which packages depend on the
        # package, which would be searching up.
        # The 'depends' subcommand corresponds to a down search.
        state["down_search"] = (args.subcmd == "depends")
        if args.package not in g:
            print(f"warning: package \"{args.package}\" not found", file=sys.stderr)
            sys.exit(1)
        elif args.dot:
            e = find_reachable_pkgs(g, args.package, exclude_pkgs=set(args.exclude), down_search=state["down_search"])
            print_graph_dot(g.subgraph(e+[args.package]))
        elif args.tree:
            tree, state = print_dep_tree(g, args.package, None, state)
            print(tree, end='')
        elif args.recursive:
            e = find_reachable_pkgs(g, args.package, exclude_pkgs=set(args.exclude), down_search=state["down_search"])
            print_pkgs(e, with_json=args.json)
        else:
            edges = g.out_edges(args.package) if state["down_search"] else g.in_edges(args.package)
            e = [i[1] for i in edges] if state["down_search"] else [i[0] for i in edges]
            print_pkgs(e, with_json=args.json)

    elif args.subcmd == 'leaves':
        if args.with_cycles:
            lv = get_leaves_plus_cycles(g)
        else:
            lv = get_leaves(g)
        if args.export:
            for p in lv:
                k = get_package_key(l, p)
                print('%s::%s=%s=%s' % (l[k].channel.channel_name, l[k].name, l[k].version, l[k].build))
        else:
            print_pkgs(lv, with_json=args.json)

    elif args.subcmd == 'deptree':
        if args.dot:
            print_graph_dot(g, exclude_pkgs=set(args.exclude))
        elif args.json:
            print_pkgs(list(g), with_json=True)
        else:
            complete_tree = ""
            for pk in get_leaves_plus_cycles(g):
                if pk not in args.exclude:
                    tree, state = print_dep_tree(g, pk, None, state)
                    complete_tree += tree
            print(''.join(complete_tree), end='')

    elif args.subcmd == 'unowned-files':
        find_unowned_files(args.prefix)

    elif args.subcmd == 'who-owns':
        find_who_owns_file(args.prefix,args.file)

    else:
        parser.print_help()
        sys.exit(1)

    #######
    # End warning messages
    #######

    # If we use a tree-based command without --full enabled
    if state["hidden_dependencies"] and not args.full:
        print(f"\n{Style.DIM}For the sake of clarity, some redundancies have been hidden.\n" +
              f"Please use the '--full' option to display them anyway.{Style.RESET_ALL}")
        if not args.small:
            print(f"\n{Style.DIM}If you are tired of seeing 'conda' and 'python' everywhere,\n" +
              f"you can use the '--small' option to hide their dependencies completely.{Style.RESET_ALL}")

    # If we use a tree-based command without --full enabled
    if state["hidden_dependencies"] and args.full:
        print(f"\n{Style.DIM}The full dependency tree shows dependencies of packages " +
              f"with cycles only once.{Style.RESET_ALL}")

if __name__ == "__main__":
    main()