Support pickling of Structs -- _encode_self instance method needs some magic

1 files changed, 25 insertions, 0 deletions

diff --git a/kafka/protocol/pickle.py b/kafka/protocol/pickle.py
new file mode 100644
index 0000000..af0d1ee
--- /dev/null
+++ b/kafka/protocol/pickle.py
@@ -0,0 +1,25 @@
+from __future__ import absolute_import
+
+import copy_reg
+import types
+
+
+def _pickle_method(method):
+    func_name = method.im_func.__name__
+    obj = method.im_self
+    cls = method.im_class
+    return _unpickle_method, (func_name, obj, cls)
+
+
+def _unpickle_method(func_name, obj, cls):
+    for cls in cls.mro():
+        try:
+            func = cls.__dict__[func_name]
+        except KeyError:
+            pass
+        else:
+            break
+        return func.__get__(obj, cls)
+
+# https://bytes.com/topic/python/answers/552476-why-cant-you-pickle-instancemethods
+copy_reg.pickle(types.MethodType, _pickle_method, _unpickle_method)