Define a pretty printer for enums

HypothesisWorks · Jan 26, 2023 · 77b0bbb · 77b0bbb
1 parent 39e3fba
commit 77b0bbb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/hypothesis-python/src/hypothesis/vendor/pretty.py b/hypothesis-python/src/hypothesis/vendor/pretty.py
@@ -762,9 +762,14 @@ def _repr_dataframe(obj, p, cycle):  # pragma: no cover
     p.break_()
 
 
+def _repr_enum(obj, p, cycle):
+    p.text(type(obj).__name__ + "." + obj.name)
+
+
 for_type_by_name("collections", "defaultdict", _defaultdict_pprint)
 for_type_by_name("collections", "OrderedDict", _ordereddict_pprint)
 for_type_by_name("ordereddict", "OrderedDict", _ordereddict_pprint)
 for_type_by_name("collections", "deque", _deque_pprint)
 for_type_by_name("collections", "Counter", _counter_pprint)
 for_type_by_name("pandas.core.frame", "DataFrame", _repr_dataframe)
+for_type_by_name("enum", "Enum", _repr_enum)
diff --git a/hypothesis-python/tests/cover/test_pretty.py b/hypothesis-python/tests/cover/test_pretty.py
@@ -50,6 +50,7 @@
 import re
 import warnings
 from collections import Counter, OrderedDict, defaultdict, deque
+from enum import Enum
 
 import pytest
 
@@ -620,3 +621,11 @@ def test_repr_call(func_name):
     assert _repr_call(func_name, (aas,), {}) == f"{fn}(\n    {aas!r},\n)"
     assert _repr_call(func_name, (), {"a": 1, "b": 2}) == f"{fn}(a=1, b=2)"
     assert _repr_call(func_name, (), {"x": aas}) == f"{fn}(\n    x={aas!r},\n)"
+
+
+class AnEnum(Enum):
+    SOME_MEMBER = 1
+
+
+def test_pretty_prints_enums_as_code():
+    assert pretty.pretty(AnEnum.SOME_MEMBER) == "AnEnum.SOME_MEMBER"