diff --git a/IPython/lib/pretty.py b/IPython/lib/pretty.py index c0553ea..83fe892 100644 --- a/IPython/lib/pretty.py +++ b/IPython/lib/pretty.py @@ -816,6 +816,51 @@ _singleton_pprinters = dict.fromkeys(map(id, [None, True, False, Ellipsis, 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 index 0c7273a..7d49e7f 100644 --- 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 test_basic_class(): 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)