From 4d715f63ec4f622e98d296e75d22d6a717b1f709 Mon Sep 17 00:00:00 2001
From: Antony Lee <anntzer.lee@gmail.com>
Date: Mon, 19 Aug 2019 00:52:30 +0200
Subject: [PATCH] Fix prettyprinting of IntEnum, Flag, and IntFlag.

---
 prettyprinter/pretty_stdlib.py   | 12 +++++++++++-
 tests/test_stdlib_definitions.py | 13 +++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/prettyprinter/pretty_stdlib.py b/prettyprinter/pretty_stdlib.py
index 2733e96..9d4eb14 100644
--- a/prettyprinter/pretty_stdlib.py
+++ b/prettyprinter/pretty_stdlib.py
@@ -303,10 +303,20 @@ def pretty_counter(counter, ctx):
     )
 
 
+# IntEnum inherits from int first, but we want Enum-style prettyprinting.
 @register_pretty('enum.Enum')
+@register_pretty('enum.IntEnum')
+@register_pretty('enum.Flag')
+@register_pretty('enum.IntFlag')
 def pretty_enum(value, ctx):
     cls = type(value)
-    return classattr(cls, value.name)
+    if value.name is not None:
+        return classattr(cls, value.name)
+    else:  # Combined flags.
+        s = str(value)
+        prefix = cls.__name__ + "."
+        assert s.startswith(prefix)
+        return classattr(cls, s[len(prefix):])
 
 
 @register_pretty('builtins.mappingproxy')
diff --git a/tests/test_stdlib_definitions.py b/tests/test_stdlib_definitions.py
index 64b8ab4..5564bbb 100644
--- a/tests/test_stdlib_definitions.py
+++ b/tests/test_stdlib_definitions.py
@@ -64,6 +64,19 @@ class TestEnum(Enum):
     expected = 'tests.test_stdlib_definitions.test_enum.<locals>.TestEnum.ONE'
     assert pformat(value) == expected
 
+    try:
+        from enum import IntFlag
+    except ImportError:
+        pass
+    else:
+        class TestFlag(IntFlag):
+            ONE = 1
+            TWO = 2
+
+        value = TestFlag.ONE | TestFlag.TWO
+        expected = 'tests.test_stdlib_definitions.test_enum.<locals>.TestFlag.TWO|ONE'
+        assert pformat(value) == expected
+
 
 def test_mappingproxytype():
     assert is_registered(