From e08720e49ae43a741f1b7cb45b4b0e52c7cc7e36 Mon Sep 17 00:00:00 2001
From: Clement Rouault <project@hakril.net>
Date: Wed, 5 Apr 2017 15:06:02 +0200
Subject: [PATCH] Fix pickle/unpikcl of Flag + StrFlag for protocol 2 +
 StrFlag.__str__ is the string + test

---
 windows/generated_def/windef.py | 10 +++++++++-
 windows/test/mytest.py          | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/windows/generated_def/windef.py b/windows/generated_def/windef.py
index b54cfd4..9116fd1 100644
--- a/windows/generated_def/windef.py
+++ b/windows/generated_def/windef.py
@@ -23,6 +23,10 @@ class Flag(long):
 
     __str__ = __repr__
 
+   # Fix pickling with protocol 2
+    def __getnewargs__(self, *args):
+        return self.name, long(self)
+
 class StrFlags(str):
     def __new__(cls, name, value):
         if isinstance(value, cls):
@@ -35,7 +39,11 @@ class StrFlags(str):
     def __repr__(self):
         return "{0}({1})".format(self.name, str.__repr__(self))
 
-    __str__ = __repr__
+    # __str__ = __repr__
+
+    # Fix pickling with protocol 2
+    def __getnewargs__(self, *args):
+        return self.name, str.__str__(self)
 
 def make_flag(name, value):
     if isinstance(value, (int, long)):
diff --git a/windows/test/mytest.py b/windows/test/mytest.py
index 0b13d8b..aa2e7aa 100644
--- a/windows/test/mytest.py
+++ b/windows/test/mytest.py
@@ -4,6 +4,7 @@ import time
 import os
 import textwrap
 import random
+import pickle
 
 from test_utils import *
 from windows.generated_def.winstructs import *
@@ -537,6 +538,29 @@ class GeneratedCodeTestCase(unittest.TestCase):
     def test_str_flags_value(self):
         self.assertEqual(windows.generated_def.MS_ENHANCED_PROV, windows.generated_def.MS_ENHANCED_PROV_A)
 
+    def _test_pickle_unpickle(self, obj, protocol=0):
+        pickled = pickle.dumps(obj, protocol)
+        unpickled = pickle.loads(pickled)
+        self.assertEqual(unpickled, obj)
+
+    def test_long_flag_picke_v0(self):
+        self._test_pickle_unpickle(windows.generated_def.PAGE_EXECUTE_READWRITE, 0)
+
+    def test_long_flag_picke_v1(self):
+        self._test_pickle_unpickle(windows.generated_def.PAGE_EXECUTE_READWRITE, 1)
+
+    def test_long_flag_picke_v2(self):
+        self._test_pickle_unpickle(windows.generated_def.PAGE_EXECUTE_READWRITE, 2)
+
+    def test_str_flag_picke_v0(self):
+        self._test_pickle_unpickle(windows.generated_def.szOID_RSA, 0)
+
+    def test_str_flag_picke_v1(self):
+        self._test_pickle_unpickle(windows.generated_def.szOID_RSA, 1)
+
+    def test_str_flag_picke_v2(self):
+        self._test_pickle_unpickle(windows.generated_def.szOID_RSA, 2)
+
 class NativeUtilsTestCase(unittest.TestCase):
     @process_64bit_only
     def test_strlenw64(self):