1from __future__ import print_function
2import os
3import sys
4from itertools import chain
5from collections import defaultdict
6import argparse
7from operator import attrgetter
8import json
9from importlib import import_module
10
11try:
12    from collections import OrderedDict
13except ImportError:
14    from ordereddict import OrderedDict
15
16pardir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
17sys.path.append(pardir)
18from pipenv.vendor.pip_shims import get_installed_distributions, FrozenRequirement
19
20import pkg_resources
21# inline:
22# from graphviz import backend, Digraph
23
24
25__version__ = '1.0.0'
26
27
28flatten = chain.from_iterable
29
30
31def build_dist_index(pkgs):
32    """Build an index pkgs by their key as a dict.
33
34    :param list pkgs: list of pkg_resources.Distribution instances
35    :returns: index of the pkgs by the pkg key
36    :rtype: dict
37
38    """
39    return dict((p.key, DistPackage(p)) for p in pkgs)
40
41
42def construct_tree(index):
43    """Construct tree representation of the pkgs from the index.
44
45    The keys of the dict representing the tree will be objects of type
46    DistPackage and the values will be list of ReqPackage objects.
47
48    :param dict index: dist index ie. index of pkgs by their keys
49    :returns: tree of pkgs and their dependencies
50    :rtype: dict
51
52    """
53    return dict((p, [ReqPackage(r, index.get(r.key))
54                     for r in p.requires()])
55                for p in index.values())
56
57
58def sorted_tree(tree):
59    """Sorts the dict representation of the tree
60
61    The root packages as well as the intermediate packages are sorted
62    in the alphabetical order of the package names.
63
64    :param dict tree: the pkg dependency tree obtained by calling
65                     `construct_tree` function
66    :returns: sorted tree
67    :rtype: collections.OrderedDict
68
69    """
70    return OrderedDict(sorted([(k, sorted(v, key=attrgetter('key')))
71                               for k, v in tree.items()],
72                              key=lambda kv: kv[0].key))
73
74
75def find_tree_root(tree, key):
76    """Find a root in a tree by it's key
77
78    :param dict tree: the pkg dependency tree obtained by calling
79                     `construct_tree` function
80    :param str key: key of the root node to find
81    :returns: a root node if found else None
82    :rtype: mixed
83
84    """
85    result = [p for p in tree.keys() if p.key == key]
86    assert len(result) in [0, 1]
87    return None if len(result) == 0 else result[0]
88
89
90def reverse_tree(tree):
91    """Reverse the dependency tree.
92
93    ie. the keys of the resulting dict are objects of type
94    ReqPackage and the values are lists of DistPackage objects.
95
96    :param dict tree: the pkg dependency tree obtained by calling
97                      `construct_tree` function
98    :returns: reversed tree
99    :rtype: dict
100
101    """
102    rtree = defaultdict(list)
103    child_keys = set(c.key for c in flatten(tree.values()))
104    for k, vs in tree.items():
105        for v in vs:
106            node = find_tree_root(rtree, v.key) or v
107            rtree[node].append(k.as_required_by(v))
108        if k.key not in child_keys:
109            rtree[k.as_requirement()] = []
110    return rtree
111
112
113def guess_version(pkg_key, default='?'):
114    """Guess the version of a pkg when pip doesn't provide it
115
116    :param str pkg_key: key of the package
117    :param str default: default version to return if unable to find
118    :returns: version
119    :rtype: string
120
121    """
122    try:
123        m = import_module(pkg_key)
124    except ImportError:
125        return default
126    else:
127        return getattr(m, '__version__', default)
128
129
130def frozen_req_from_dist(dist):
131    try:
132        return FrozenRequirement.from_dist(dist)
133    except TypeError:
134        return FrozenRequirement.from_dist(dist, [])
135
136
137class Package(object):
138    """Abstract class for wrappers around objects that pip returns.
139
140    This class needs to be subclassed with implementations for
141    `render_as_root` and `render_as_branch` methods.
142
143    """
144
145    def __init__(self, obj):
146        self._obj = obj
147        self.project_name = obj.project_name
148        self.key = obj.key
149
150    def render_as_root(self, frozen):
151        return NotImplementedError
152
153    def render_as_branch(self, frozen):
154        return NotImplementedError
155
156    def render(self, parent=None, frozen=False):
157        if not parent:
158            return self.render_as_root(frozen)
159        else:
160            return self.render_as_branch(frozen)
161
162    @staticmethod
163    def frozen_repr(obj):
164        fr = frozen_req_from_dist(obj)
165        return str(fr).strip()
166
167    def __getattr__(self, key):
168        return getattr(self._obj, key)
169
170    def __repr__(self):
171        return '<{0}("{1}")>'.format(self.__class__.__name__, self.key)
172
173
174class DistPackage(Package):
175    """Wrapper class for pkg_resources.Distribution instances
176
177      :param obj: pkg_resources.Distribution to wrap over
178      :param req: optional ReqPackage object to associate this
179                  DistPackage with. This is useful for displaying the
180                  tree in reverse
181    """
182
183    def __init__(self, obj, req=None):
184        super(DistPackage, self).__init__(obj)
185        self.version_spec = None
186        self.req = req
187
188    def render_as_root(self, frozen):
189        if not frozen:
190            return '{0}=={1}'.format(self.project_name, self.version)
191        else:
192            return self.__class__.frozen_repr(self._obj)
193
194    def render_as_branch(self, frozen):
195        assert self.req is not None
196        if not frozen:
197            parent_ver_spec = self.req.version_spec
198            parent_str = self.req.project_name
199            if parent_ver_spec:
200                parent_str += parent_ver_spec
201            return (
202                '{0}=={1} [requires: {2}]'
203            ).format(self.project_name, self.version, parent_str)
204        else:
205            return self.render_as_root(frozen)
206
207    def as_requirement(self):
208        """Return a ReqPackage representation of this DistPackage"""
209        return ReqPackage(self._obj.as_requirement(), dist=self)
210
211    def as_required_by(self, req):
212        """Return a DistPackage instance associated to a requirement
213
214        This association is necessary for displaying the tree in
215        reverse.
216
217        :param ReqPackage req: the requirement to associate with
218        :returns: DistPackage instance
219
220        """
221        return self.__class__(self._obj, req)
222
223    def as_dict(self):
224        return {'key': self.key,
225                'package_name': self.project_name,
226                'installed_version': self.version}
227
228
229class ReqPackage(Package):
230    """Wrapper class for Requirements instance
231
232      :param obj: The `Requirements` instance to wrap over
233      :param dist: optional `pkg_resources.Distribution` instance for
234                   this requirement
235    """
236
237    UNKNOWN_VERSION = '?'
238
239    def __init__(self, obj, dist=None):
240        super(ReqPackage, self).__init__(obj)
241        self.dist = dist
242
243    @property
244    def version_spec(self):
245        specs = sorted(self._obj.specs, reverse=True)  # `reverse` makes '>' prior to '<'
246        return ','.join([''.join(sp) for sp in specs]) if specs else None
247
248    @property
249    def installed_version(self):
250        if not self.dist:
251            return guess_version(self.key, self.UNKNOWN_VERSION)
252        return self.dist.version
253
254    def is_conflicting(self):
255        """If installed version conflicts with required version"""
256        # unknown installed version is also considered conflicting
257        if self.installed_version == self.UNKNOWN_VERSION:
258            return True
259        ver_spec = (self.version_spec if self.version_spec else '')
260        req_version_str = '{0}{1}'.format(self.project_name, ver_spec)
261        req_obj = pkg_resources.Requirement.parse(req_version_str)
262        return self.installed_version not in req_obj
263
264    def render_as_root(self, frozen):
265        if not frozen:
266            return '{0}=={1}'.format(self.project_name, self.installed_version)
267        elif self.dist:
268            return self.__class__.frozen_repr(self.dist._obj)
269        else:
270            return self.project_name
271
272    def render_as_branch(self, frozen):
273        if not frozen:
274            req_ver = self.version_spec if self.version_spec else 'Any'
275            return (
276                '{0} [required: {1}, installed: {2}]'
277                ).format(self.project_name, req_ver, self.installed_version)
278        else:
279            return self.render_as_root(frozen)
280
281    def as_dict(self):
282        return {'key': self.key,
283                'package_name': self.project_name,
284                'installed_version': self.installed_version,
285                'required_version': self.version_spec}
286
287
288def render_tree(tree, list_all=True, show_only=None, frozen=False, exclude=None):
289    """Convert tree to string representation
290
291    :param dict tree: the package tree
292    :param bool list_all: whether to list all the pgks at the root
293                          level or only those that are the
294                          sub-dependencies
295    :param set show_only: set of select packages to be shown in the
296                          output. This is optional arg, default: None.
297    :param bool frozen: whether or not show the names of the pkgs in
298                        the output that's favourable to pip --freeze
299    :param set exclude: set of select packages to be excluded from the
300                          output. This is optional arg, default: None.
301    :returns: string representation of the tree
302    :rtype: str
303
304    """
305    tree = sorted_tree(tree)
306    branch_keys = set(r.key for r in flatten(tree.values()))
307    nodes = tree.keys()
308    use_bullets = not frozen
309
310    key_tree = dict((k.key, v) for k, v in tree.items())
311    get_children = lambda n: key_tree.get(n.key, [])
312
313    if show_only:
314        nodes = [p for p in nodes
315                 if p.key in show_only or p.project_name in show_only]
316    elif not list_all:
317        nodes = [p for p in nodes if p.key not in branch_keys]
318
319    def aux(node, parent=None, indent=0, chain=None):
320        if exclude and (node.key in exclude or node.project_name in exclude):
321            return []
322        if chain is None:
323            chain = [node.project_name]
324        node_str = node.render(parent, frozen)
325        if parent:
326            prefix = ' '*indent + ('- ' if use_bullets else '')
327            node_str = prefix + node_str
328        result = [node_str]
329        children = [aux(c, node, indent=indent+2,
330                        chain=chain+[c.project_name])
331                    for c in get_children(node)
332                    if c.project_name not in chain]
333        result += list(flatten(children))
334        return result
335
336    lines = flatten([aux(p) for p in nodes])
337    return '\n'.join(lines)
338
339
340def render_json(tree, indent):
341    """Converts the tree into a flat json representation.
342
343    The json repr will be a list of hashes, each hash having 2 fields:
344      - package
345      - dependencies: list of dependencies
346
347    :param dict tree: dependency tree
348    :param int indent: no. of spaces to indent json
349    :returns: json representation of the tree
350    :rtype: str
351
352    """
353    return json.dumps([{'package': k.as_dict(),
354                        'dependencies': [v.as_dict() for v in vs]}
355                       for k, vs in tree.items()],
356                      indent=indent)
357
358
359def render_json_tree(tree, indent):
360    """Converts the tree into a nested json representation.
361
362    The json repr will be a list of hashes, each hash having the following fields:
363      - package_name
364      - key
365      - required_version
366      - installed_version
367      - dependencies: list of dependencies
368
369    :param dict tree: dependency tree
370    :param int indent: no. of spaces to indent json
371    :returns: json representation of the tree
372    :rtype: str
373
374    """
375    tree = sorted_tree(tree)
376    branch_keys = set(r.key for r in flatten(tree.values()))
377    nodes = [p for p in tree.keys() if p.key not in branch_keys]
378    key_tree = dict((k.key, v) for k, v in tree.items())
379    get_children = lambda n: key_tree.get(n.key, [])
380
381    def aux(node, parent=None, chain=None):
382        if chain is None:
383            chain = [node.project_name]
384
385        d = node.as_dict()
386        if parent:
387            d['required_version'] = node.version_spec if node.version_spec else 'Any'
388        else:
389            d['required_version'] = d['installed_version']
390
391        d['dependencies'] = [
392            aux(c, parent=node, chain=chain+[c.project_name])
393            for c in get_children(node)
394            if c.project_name not in chain
395        ]
396
397        return d
398
399    return json.dumps([aux(p) for p in nodes], indent=indent)
400
401
402def dump_graphviz(tree, output_format='dot'):
403    """Output dependency graph as one of the supported GraphViz output formats.
404
405    :param dict tree: dependency graph
406    :param string output_format: output format
407    :returns: representation of tree in the specified output format
408    :rtype: str or binary representation depending on the output format
409
410    """
411    try:
412        from graphviz import backend, Digraph
413    except ImportError:
414        print('graphviz is not available, but necessary for the output '
415              'option. Please install it.', file=sys.stderr)
416        sys.exit(1)
417
418    if output_format not in backend.FORMATS:
419        print('{0} is not a supported output format.'.format(output_format),
420              file=sys.stderr)
421        print('Supported formats are: {0}'.format(
422            ', '.join(sorted(backend.FORMATS))), file=sys.stderr)
423        sys.exit(1)
424
425    graph = Digraph(format=output_format)
426    for package, deps in tree.items():
427        project_name = package.project_name
428        label = '{0}\n{1}'.format(project_name, package.version)
429        graph.node(project_name, label=label)
430        for dep in deps:
431            label = dep.version_spec
432            if not label:
433                label = 'any'
434            graph.edge(project_name, dep.project_name, label=label)
435
436    # Allow output of dot format, even if GraphViz isn't installed.
437    if output_format == 'dot':
438        return graph.source
439
440    # As it's unknown if the selected output format is binary or not, try to
441    # decode it as UTF8 and only print it out in binary if that's not possible.
442    try:
443        return graph.pipe().decode('utf-8')
444    except UnicodeDecodeError:
445        return graph.pipe()
446
447
448def print_graphviz(dump_output):
449    """Dump the data generated by GraphViz to stdout.
450
451    :param dump_output: The output from dump_graphviz
452    """
453    if hasattr(dump_output, 'encode'):
454        print(dump_output)
455    else:
456        with os.fdopen(sys.stdout.fileno(), 'wb') as bytestream:
457            bytestream.write(dump_output)
458
459
460def conflicting_deps(tree):
461    """Returns dependencies which are not present or conflict with the
462    requirements of other packages.
463
464    e.g. will warn if pkg1 requires pkg2==2.0 and pkg2==1.0 is installed
465
466    :param tree: the requirements tree (dict)
467    :returns: dict of DistPackage -> list of unsatisfied/unknown ReqPackage
468    :rtype: dict
469
470    """
471    conflicting = defaultdict(list)
472    for p, rs in tree.items():
473        for req in rs:
474            if req.is_conflicting():
475                conflicting[p].append(req)
476    return conflicting
477
478
479def cyclic_deps(tree):
480    """Return cyclic dependencies as list of tuples
481
482    :param list pkgs: pkg_resources.Distribution instances
483    :param dict pkg_index: mapping of pkgs with their respective keys
484    :returns: list of tuples representing cyclic dependencies
485    :rtype: generator
486
487    """
488    key_tree = dict((k.key, v) for k, v in tree.items())
489    get_children = lambda n: key_tree.get(n.key, [])
490    cyclic = []
491    for p, rs in tree.items():
492        for req in rs:
493            if p.key in map(attrgetter('key'), get_children(req)):
494                cyclic.append((p, req, p))
495    return cyclic
496
497
498def get_parser():
499    parser = argparse.ArgumentParser(description=(
500        'Dependency tree of the installed python packages'
501    ))
502    parser.add_argument('-v', '--version', action='version',
503                        version='{0}'.format(__version__))
504    parser.add_argument('-f', '--freeze', action='store_true',
505                        help='Print names so as to write freeze files')
506    parser.add_argument('-a', '--all', action='store_true',
507                        help='list all deps at top level')
508    parser.add_argument('-l', '--local-only',
509                        action='store_true', help=(
510                            'If in a virtualenv that has global access '
511                            'do not show globally installed packages'
512                        ))
513    parser.add_argument('-u', '--user-only', action='store_true',
514                        help=(
515                            'Only show installations in the user site dir'
516                        ))
517    parser.add_argument('-w', '--warn', action='store', dest='warn',
518                        nargs='?', default='suppress',
519                        choices=('silence', 'suppress', 'fail'),
520                        help=(
521                            'Warning control. "suppress" will show warnings '
522                            'but return 0 whether or not they are present. '
523                            '"silence" will not show warnings at all and '
524                            'always return 0. "fail" will show warnings and '
525                            'return 1 if any are present. The default is '
526                            '"suppress".'
527                        ))
528    parser.add_argument('-r', '--reverse', action='store_true',
529                        default=False, help=(
530                            'Shows the dependency tree in the reverse fashion '
531                            'ie. the sub-dependencies are listed with the '
532                            'list of packages that need them under them.'
533                        ))
534    parser.add_argument('-p', '--packages',
535                        help=(
536                            'Comma separated list of select packages to show '
537                            'in the output. If set, --all will be ignored.'
538                        ))
539    parser.add_argument('-e', '--exclude',
540                        help=(
541                            'Comma separated list of select packages to exclude '
542                            'from the output. If set, --all will be ignored.'
543                        ), metavar='PACKAGES')
544    parser.add_argument('-j', '--json', action='store_true', default=False,
545                        help=(
546                            'Display dependency tree as json. This will yield '
547                            '"raw" output that may be used by external tools. '
548                            'This option overrides all other options.'
549                        ))
550    parser.add_argument('--json-tree', action='store_true', default=False,
551                        help=(
552                            'Display dependency tree as json which is nested '
553                            'the same way as the plain text output printed by default. '
554                            'This option overrides all other options (except --json).'
555                        ))
556    parser.add_argument('--graph-output', dest='output_format',
557                        help=(
558                            'Print a dependency graph in the specified output '
559                            'format. Available are all formats supported by '
560                            'GraphViz, e.g.: dot, jpeg, pdf, png, svg'
561                        ))
562    return parser
563
564
565def _get_args():
566    parser = get_parser()
567    return parser.parse_args()
568
569
570def main():
571    args = _get_args()
572    pkgs = get_installed_distributions(local_only=args.local_only,
573                                       user_only=args.user_only)
574
575    dist_index = build_dist_index(pkgs)
576    tree = construct_tree(dist_index)
577
578    if args.json:
579        print(render_json(tree, indent=4))
580        return 0
581    elif args.json_tree:
582        print(render_json_tree(tree, indent=4))
583        return 0
584    elif args.output_format:
585        output = dump_graphviz(tree, output_format=args.output_format)
586        print_graphviz(output)
587        return 0
588
589    return_code = 0
590
591    # show warnings about possibly conflicting deps if found and
592    # warnings are enabled
593    if args.warn != 'silence':
594        conflicting = conflicting_deps(tree)
595        if conflicting:
596            print('Warning!!! Possibly conflicting dependencies found:',
597                  file=sys.stderr)
598            for p, reqs in conflicting.items():
599                pkg = p.render_as_root(False)
600                print('* {}'.format(pkg), file=sys.stderr)
601                for req in reqs:
602                    req_str = req.render_as_branch(False)
603                    print(' - {}'.format(req_str), file=sys.stderr)
604            print('-'*72, file=sys.stderr)
605
606        cyclic = cyclic_deps(tree)
607        if cyclic:
608            print('Warning!! Cyclic dependencies found:', file=sys.stderr)
609            for a, b, c in cyclic:
610                print('* {0} => {1} => {2}'.format(a.project_name,
611                                                   b.project_name,
612                                                   c.project_name),
613                      file=sys.stderr)
614            print('-'*72, file=sys.stderr)
615
616        if args.warn == 'fail' and (conflicting or cyclic):
617            return_code = 1
618
619    show_only = set(args.packages.split(',')) if args.packages else None
620    exclude = set(args.exclude.split(',')) if args.exclude else None
621
622    if show_only and exclude and (show_only & exclude):
623        print('Conflicting packages found in --packages and --exclude lists.', file=sys.stderr)
624        sys.exit(1)
625
626    tree = render_tree(tree if not args.reverse else reverse_tree(tree),
627                       list_all=args.all, show_only=show_only,
628                       frozen=args.freeze, exclude=exclude)
629    print(tree)
630    return return_code
631
632
633if __name__ == '__main__':
634    sys.exit(main())
635