Commits

Łukasz Langa  committed 9e33d17

Support for sets and frozensets introduced.

  • Participants
  • Parent commits cc3eec8

Comments (0)

Files changed (3)

 - Port checks for recursive collections from reprlib or pprint.
 - Add safe string breaking.
-- Add support for: dict, namedtuple and sets.
+- Add support for: dict and namedtuple.
 - Support for specifying custom formatters in select_formatter() and pprint().

File nattyprint.py

     LAST_SEPARATOR = 'if 1 element' # possible values: 'never',
                                     #                  'if 1 element',
                                     #                  'always'
+    # A symbol for the container when it holds no elements
+    EMPTY_CONTAINER = '()'
+    # Character used as the prefix for a container with at least 1 element
     PREFIX = '('
+    # Character used for separation of elements 
     SEPARATOR = ','
+    # Character used as the suffix for a container with at least 1 element
     SUFFIX = ')'
 
     @classmethod
         return other.__class__ is tuple
 
     @classmethod
+    def _pprint_generate_elements(cls, instance):
+        for elem in instance:
+            yield elem
+
+    @classmethod
     def _pprint(cls, instance, width, height, buffer):
         def enough_space(elem_width=1, elem_height=1):
             closing_width = max(len(cls.SEPARATOR), len(cls.SUFFIX))
             # the -1 above is the yet unprinted space after the separator
             return cls.MULTIPLE_VALUES and result
 
+        if len(instance) == 0:
+            buffer.write(cls.EMPTY_CONTAINER)
+            return
+
         buffer.write(cls.PREFIX)
         available_width = sys.maxsize
         first_elem = True
-        for elem in instance:
+        for elem in cls._pprint_generate_elements(instance):
             got_newline = False
             if not first_elem:
                 if enough_space():
     LAST_SEPARATOR = 'never'
     PREFIX = '['
     SUFFIX = ']'
+    EMPTY_CONTAINER = '[]'
 
     @classmethod
     def __accepts__(cls, other):
     pass
 
 
+class SetFormatter(TupleFormatter):
+    MULTIPLE_VALUES = True
+    LAST_SEPARATOR = 'never'
+    PREFIX = '{'
+    SUFFIX = '}'
+    EMPTY_CONTAINER = 'set()'
+
+    @classmethod
+    def _pprint_generate_elements(cls, instance):
+        try:
+            for elem in sorted(instance):
+                yield elem
+        except TypeError:
+            for elem in list(instance):
+                yield elem
+
+
+class FrozenSetFormatter(SetFormatter):
+    PREFIX = 'frozenset(('
+    SUFFIX = '))'
+    EMPTY_CONTAINER = 'frozenset()'
+
+
 formatters = []
 
 _builtin_formatters = {
     bytes: BytesFormatter,
     dict: DictFormatter,
+    frozenset: FrozenSetFormatter,
     list: ListFormatter,
+    set: SetFormatter,
     str: StringFormatter,
     tuple: TupleFormatter,
 }
     pprint.pprint."""
 
     if indent is not None:
-        raise NotImplementedError("Not implemented: changing indentation.")
+        raise NotImplementedError("Not implemented: changing indentation "
+                                  "amount.")
     if depth is not None:
         raise NotImplementedError("Not implemented: changing depth.")
     formatter = select_formatter(object)

File test_nattyprint.py

               expected=expected_3_not_first)
 
 class TestFormatter(unittest.TestCase):
-    def check(self, formatter, object, width=None, expected=None):
+    def check(self, formatter, object, width=None, expected=None,
+              expected_in=None):
         """Expected is a three-tuple of (text, width, height)."""
-        self.assertEqual(formatter.__pprint__(object, width=width), expected)
+        got = formatter.__pprint__(object, width=width)
+        if expected is not None:
+            self.assertEqual(got, expected)
+        elif expected_in is not None and len(expected_in) > 0:
+            for exp in expected_in:
+                if exp == got:
+                    return
+            self.assertFalse("Got: '{}', expected one of: {}".format(
+                             got, expected_in))
+
 
     def test_simple_values(self):
         f = nattyprint.Formatter()
         """).strip()
         check(x, width=44, expected=(expected_repr, 11, 6))
 
+    def test_simple_sets(self):
+        f = nattyprint.SetFormatter()
+        check = functools.partial(self.check, f)
+        check(set(), expected=("set()", 5, 1))
+        check({5, 1, 4, 2, 3}, expected=("{1, 2, 3, 4, 5}", 15, 1))
+        check({5, 1, 4, 2, 3}, width=8, expected=("{1, 2,\n 3, 4,\n 5}", 3, 3))
+        check(set("eyuioa"), width=8, expected=(dedent("""
+              {'a',
+               'e',
+               'i',
+               'o',
+               'u',
+               'y'}
+              """).strip(), 5, 6))
+        check({1, "one", None}, width=8, expected_in=(
+                ("{1,\n 'one',\n None}", 6, 3),
+                ("{'one',\n 1,\n None}", 6, 3),
+                ("{1,\n None,\n 'one'}", 7, 3),
+                ("{None,\n 1,\n 'one'}", 7, 3),
+                ("{None,\n 'one',\n 1}", 3, 3),
+                ("{'one',\n None,\n 1}", 3, 3),
+                ))
+
+    def test_simple_frozensets(self):
+        f = nattyprint.FrozenSetFormatter()
+        check = functools.partial(self.check, f)
+        check(frozenset(), expected=("frozenset()", 11, 1))
+        check(frozenset((5, 1, 4, 2, 3)), expected=(
+              "frozenset((1, 2, 3, 4, 5))", 26, 1))
+        check(frozenset((5, 1, 4, 2, 3)), width=16, expected=(
+              dedent("""
+              frozenset((1,
+                         2,
+                         3,
+                         4,
+                         5))
+              """).strip(), 14, 5))
+        check(frozenset("eyuioa"), width=16, expected=(dedent("""
+              frozenset(('a',
+                         'e',
+                         'i',
+                         'o',
+                         'u',
+                         'y'))
+              """).strip(), 16, 6))
+        check({1, "one", None}, width=16, expected_in=(
+                (dedent("""
+                frozenset((1,
+                           'one',
+                           None))
+                """).strip(), 17, 3),
+                (dedent("""
+                frozenset(('one',
+                           1,
+                           None))
+                """).strip(), 17, 3),
+                (dedent("""
+                frozenset((1,
+                           None,
+                           'one'))
+                """).strip(), 18, 3),
+                (dedent("""
+                frozenset((None,
+                           1,
+                           'one'))
+                """).strip(), 18, 3),
+                (dedent("""
+                frozenset((None,
+                           'one',
+                           1))
+                """).strip(), 14, 3),
+                (dedent("""
+                frozenset(('one',
+                           None,
+                           1))
+                """).strip(), 14, 3),
+                ))
+
 if __name__ == '__main__':
     unittest.main()