Fix pickle/unpikcl of Flag + StrFlag for protocol 2 + StrFlag.__str__ is the string + test

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)):

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):