Add pprint support for containers from collections

RazerM · RazerM · commit 533e8c3a63af · 2015-05-23T03:29:23.000+02:00
for_type_by_name used to keep compatiblity with Python 2.4 rather than catching ImportError.

deque is already imported, but uses same mechanism for consistency.

Tests added for each type.
diff --git a/IPython/lib/pretty.py b/IPython/lib/pretty.py
@@ -816,6 +816,51 @@ def for_type_by_name(type_module, type_name, func):
                                       NotImplemented]), _repr_pprint)
 
 
+def _defaultdict_pprint(obj, p, cycle):
+    name = 'defaultdict'
+    p.begin_group(len(name) + 1, name + '(')
+    if cycle:
+        p.text('...')
+    else:
+        p.pretty(obj.default_factory)
+        p.text(',')
+        p.breakable()
+        p.pretty(dict(obj))
+    p.end_group(len(name) + 1, ')')
+
+def _ordereddict_pprint(obj, p, cycle):
+    name = 'OrderedDict'
+    p.begin_group(len(name) + 1, name + '(')
+    if cycle:
+        p.text('...')
+    elif len(obj):
+            p.pretty(list(obj.items()))
+    p.end_group(len(name) + 1, ')')
+
+def _deque_pprint(obj, p, cycle):
+    name = 'deque'
+    p.begin_group(len(name) + 1, name + '(')
+    if cycle:
+        p.text('...')
+    else:
+        p.pretty(list(obj))
+    p.end_group(len(name) + 1, ')')
+
+
+def _counter_pprint(obj, p, cycle):
+    name = 'Counter'
+    p.begin_group(len(name) + 1, name + '(')
+    if cycle:
+        p.text('...')
+    elif len(obj):
+            p.pretty(dict(obj))
+    p.end_group(len(name) + 1, ')')
+
+for_type_by_name('collections', 'defaultdict', _defaultdict_pprint)
+for_type_by_name('collections', 'OrderedDict', _ordereddict_pprint)
+for_type_by_name('collections', 'deque', _deque_pprint)
+for_type_by_name('collections', 'Counter', _counter_pprint)
+
 if __name__ == '__main__':
     from random import randrange
     class Foo(object):
diff --git a/IPython/lib/tests/test_pretty.py b/IPython/lib/tests/test_pretty.py
@@ -6,6 +6,8 @@
 
 from __future__ import print_function
 
+from collections import Counter, defaultdict, deque, OrderedDict
+
 import nose.tools as nt
 
 from IPython.lib import pretty
@@ -268,3 +270,89 @@ def type_pprint_wrapper(obj, p, cycle):
 
     nt.assert_equal(output, '%s.MyObj' % __name__)
     nt.assert_true(type_pprint_wrapper.called)
+
+
+def test_collections_defaultdict():
+    # Create defaultdicts with cycles
+    a = defaultdict()
+    a.default_factory = a
+    b = defaultdict(list)
+    b['key'] = b
+
+    # Dictionary order cannot be relied on, test against single keys.
+    cases = [
+        (defaultdict(list), 'defaultdict(list, {})'),
+        (defaultdict(list, {'key': '-' * 50}),
+         "defaultdict(list,\n"
+         "            {'key': '--------------------------------------------------'})"),
+        (a, 'defaultdict(defaultdict(...), {})'),
+        (b, "defaultdict(list, {'key': defaultdict(...)})"),
+    ]
+    for obj, expected in cases:
+        nt.assert_equal(pretty.pretty(obj), expected)
+
+
+def test_collections_ordereddict():
+    # Create OrderedDict with cycle
+    a = OrderedDict()
+    a['key'] = a
+
+    cases = [
+        (OrderedDict(), 'OrderedDict()'),
+        (OrderedDict((i, i) for i in range(1000, 1010)),
+         'OrderedDict([(1000, 1000),\n'
+         '             (1001, 1001),\n'
+         '             (1002, 1002),\n'
+         '             (1003, 1003),\n'
+         '             (1004, 1004),\n'
+         '             (1005, 1005),\n'
+         '             (1006, 1006),\n'
+         '             (1007, 1007),\n'
+         '             (1008, 1008),\n'
+         '             (1009, 1009)])'),
+        (a, "OrderedDict([('key', OrderedDict(...))])"),
+    ]
+    for obj, expected in cases:
+        nt.assert_equal(pretty.pretty(obj), expected)
+
+
+def test_collections_deque():
+    # Create deque with cycle
+    a = deque()
+    a.append(a)
+
+    cases = [
+        (deque(), 'deque([])'),
+        (deque(i for i in range(1000, 1020)),
+         'deque([1000,\n'
+         '       1001,\n'
+         '       1002,\n'
+         '       1003,\n'
+         '       1004,\n'
+         '       1005,\n'
+         '       1006,\n'
+         '       1007,\n'
+         '       1008,\n'
+         '       1009,\n'
+         '       1010,\n'
+         '       1011,\n'
+         '       1012,\n'
+         '       1013,\n'
+         '       1014,\n'
+         '       1015,\n'
+         '       1016,\n'
+         '       1017,\n'
+         '       1018,\n'
+         '       1019])'),
+        (a, 'deque([deque(...)])'),
+    ]
+    for obj, expected in cases:
+        nt.assert_equal(pretty.pretty(obj), expected)
+
+def test_collections_counter():
+    cases = [
+        (Counter(), 'Counter()'),
+        (Counter(a=1), "Counter({'a': 1})"),
+    ]
+    for obj, expected in cases:
+        nt.assert_equal(pretty.pretty(obj), expected)
