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