add OrderedDict and frozenset support

Author: samuelcolvin

devtools/prettier.py

@@ -17,6 +17,7 @@
 PARENTHESES_LOOKUP = [
     (list, '[', ']'),
     (set, '{', '}'),
+    (frozenset, 'frozenset({', '})'),
 ]
 __all__ = ['PrettyFormat', 'pformat', 'pprint']
 
@@ -37,9 +38,9 @@ def __init__(self,
         self._width = width
         self._type_lookup = [
             (dict, self._format_dict),
-            ((tuple, list, set), self._format_list_like),
             (str, self._format_str),
             (bytes, self._format_bytes),
+            ((tuple, list, set, frozenset), self._format_list_like),
             (collections.Generator, self._format_generators),
         ]
 
@@ -68,13 +69,18 @@ def _format(self, value: Any, indent_current: int, indent_first: bool):
             self._format_raw(value, value_repr, indent_current, indent_new)
 
     def _format_dict(self, value: dict, value_repr: str, indent_current: int, indent_new: int):
-        self._stream.write('{\n')
+        open_, before_, split_, after_, close_ = '{\n', indent_new * self._c, ': ', ',\n', '}'
+        if isinstance(value, collections.OrderedDict):
+            open_, split_, after_, close_ = 'OrderedDict([\n', ', ', '),\n', '])'
+            before_ += '('
+        self._stream.write(open_)
         for k, v in value.items():
-            self._format(k, indent_new, True)
-            self._stream.write(': ')
+            self._stream.write(before_)
+            self._format(k, indent_new, False)
+            self._stream.write(split_)
             self._format(v, indent_new, False)
-            self._stream.write(',\n')
-        self._stream.write(indent_current * self._c + '}')
+            self._stream.write(after_)
+        self._stream.write(indent_current * self._c + close_)
 
     def _format_list_like(self, value: Union[list, tuple, set], value_repr: str, indent_current: int, indent_new: int):
         open_, close_ = '(', ')'

tests/test_prettier.py

@@ -1,4 +1,5 @@
 import string
+from collections import OrderedDict
 
 import numpy
 
@@ -135,3 +136,25 @@ def test_indent_numpy_short():
     assert v == """{
     'numpy test': array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
 }"""
+
+
+def test_ordered_dict():
+    v = pformat(OrderedDict([(1, 2), (3, 4), (5, 6)]))
+    print(v)
+    assert v == """\
+OrderedDict([
+    (1, 2),
+    (3, 4),
+    (5, 6),
+])"""
+
+
+def test_frozenset():
+    v = pformat(frozenset(range(3)))
+    print(v)
+    assert v == """\
+frozenset({
+    0,
+    1,
+    2,
+})"""