correct labels for dicts

Fixes #49

Author: graingert

objgraph.py

@@ -72,6 +72,12 @@
     # Python 3.x compatibility
     iteritems = dict.items
 
+try:
+    zip_longest = itertools.izip_longest
+except AttributeError:  # pragma: PY3
+    # Python 3.x compatibility
+    zip_longest = itertools.zip_longest
+
 IS_INTERACTIVE = False
 try:  # pragma: nocover
     import graphviz
@@ -1000,7 +1006,10 @@ def _show_graph(objs, edge_func, swap_source_target,
             continue
         if cull_func is not None and cull_func(target):
             continue
-        neighbours = edge_func(target)
+        edges = edge_func(target)
+        counts = collections.Counter(id(v) for v in edges)
+        neighbours = list({id(v): v for v in edges}.values())
+        del edges
         ignore.add(id(neighbours))
         n = 0
         skipped = 0
@@ -1016,9 +1025,13 @@ def _show_graph(objs, edge_func, swap_source_target,
                 srcnode, tgtnode = target, source
             else:
                 srcnode, tgtnode = source, target
-            elabel = _edge_label(srcnode, tgtnode, shortnames)
-            f.write('  %s -> %s%s;\n' % (_obj_node_id(srcnode),
-                                         _obj_node_id(tgtnode), elabel))
+            for elabel, _ in zip_longest(
+                _edge_labels(srcnode, tgtnode, shortnames),
+                range(counts[id(source)]),
+                fillvalue='',
+            ):
+                f.write('  %s -> %s%s;\n' % (_obj_node_id(srcnode),
+                                             _obj_node_id(tgtnode), elabel))
             if id(source) not in depth:
                 depth[id(source)] = tdepth + 1
                 queue.append(source)
@@ -1208,43 +1221,51 @@ def _gradient(start_color, end_color, depth, max_depth):
     return h, s, v
 
 
-def _edge_label(source, target, shortnames=True):
+def _edge_labels(source, target, shortnames=True):
     if (_isinstance(target, dict)
             and target is getattr(source, '__dict__', None)):
-        return ' [label="__dict__",weight=10]'
+        return [' [label="__dict__",weight=10]']
     if _isinstance(source, types.FrameType):
         if target is source.f_locals:
-            return ' [label="f_locals",weight=10]'
+            return [' [label="f_locals",weight=10]']
         if target is source.f_globals:
-            return ' [label="f_globals",weight=10]'
+            return [' [label="f_globals",weight=10]']
     if _isinstance(source, types.MethodType):
         try:
             if target is source.__self__:
-                return ' [label="__self__",weight=10]'
+                return [' [label="__self__",weight=10]']
             if target is source.__func__:
-                return ' [label="__func__",weight=10]'
+                return [' [label="__func__",weight=10]']
         except AttributeError:  # pragma: nocover
             # Python < 2.6 compatibility
             if target is source.im_self:
-                return ' [label="im_self",weight=10]'
+                return [' [label="im_self",weight=10]']
             if target is source.im_func:
-                return ' [label="im_func",weight=10]'
+                return [' [label="im_func",weight=10]']
     if _isinstance(source, types.FunctionType):
-        for k in dir(source):
-            if target is getattr(source, k):
-                return ' [label="%s",weight=10]' % _quote(k)
+        return [
+            ' [label="%s",weight=10]' % _quote(k)
+            for k in dir(source)
+            if target is getattr(source, k)
+        ]
     if _isinstance(source, dict):
-        for k, v in iteritems(source):
-            if v is target:
-                if _isinstance(k, basestring) and _is_identifier(k):
-                    return ' [label="%s",weight=2]' % _quote(k)
-                else:
-                    if shortnames:
-                        tn = _short_typename(k)
-                    else:
-                        tn = _long_typename(k)
-                    return ' [label="%s"]' % _quote(tn + "\n" + _safe_repr(k))
-    return ''
+        tn = _short_typename if shortnames else _long_typename
+        return [
+            (
+                ' [label="%s",weight=2]' % _quote(k)
+                if _isinstance(k, basestring) and _is_identifier(k)
+                else (
+                    ' [label="%s"]' % _quote(tn(k) + "\n" + _safe_repr(k))
+                )
+            )
+            for k, v in iteritems(source)
+            if v is target
+        ]
+    return []
+
+
+def _edge_label(*args, **kwargs):
+    return next(iter(_edge_labels(*args, **kwargs)), '')
 
 
 _is_identifier = re.compile('[a-zA-Z_][a-zA-Z_0-9]*$').match

tests.py

@@ -1,4 +1,5 @@
 #!/usr/bin/python
+import collections
 import doctest
 import gc
 import glob
@@ -247,6 +248,40 @@ def test_cull_func(self):
                                 label_a=label_a,
                                 label_b=label_b))
 
+    @skipIf(
+        sys.version_info < (3, 6),
+        "Python < 3.6 dicts have random iteration order",
+    )
+    def test_dict(self):
+        d = dict.fromkeys("abcdefg")
+        output = StringIO()
+        objgraph.show_refs(d, output=output)
+        self.assertEqual(
+            output.getvalue(),
+            textwrap.dedent(
+                """\
+                digraph ObjectGraph {{
+                  node[shape=box, style=filled, fillcolor=white];
+                  {d_id}[fontcolor=red];
+                  {d_id}[label="dict\\n7 items"];
+                  {d_id}[fillcolor="0,0,1"];
+                  {d_id} -> {none_id} [label="a",weight=2];
+                  {d_id} -> {none_id} [label="b",weight=2];
+                  {d_id} -> {none_id} [label="c",weight=2];
+                  {d_id} -> {none_id} [label="d",weight=2];
+                  {d_id} -> {none_id} [label="e",weight=2];
+                  {d_id} -> {none_id} [label="f",weight=2];
+                  {d_id} -> {none_id} [label="g",weight=2];
+                  {none_id}[label="NoneType\\nNone"];
+                  {none_id}[fillcolor="0,0,0.766667"];
+                }}
+                """
+            ).format(
+                d_id=objgraph._obj_node_id(d),
+                none_id=objgraph._obj_node_id(None),
+            ),
+        )
+
     @mock.patch('objgraph.IS_INTERACTIVE', True)
     @mock.patch('objgraph.graphviz', create=True)
     def test_ipython(self, mock_graphviz):