Merge pull request #118 from HyukjinKwon/pickle_module_name

Sets the missing module name when pickling a function

Author: mariusvniekerk

cloudpickle/cloudpickle.py

@@ -508,6 +508,7 @@ def save_function_tuple(self, func):
         save(f_globals)
         save(defaults)
         save(dct)
+        save(func.__module__)
         save(closure_values)
         write(pickle.TUPLE)
         write(pickle.REDUCE)  # applies _fill_function on the tuple
@@ -1001,13 +1002,14 @@ def __reduce__(cls):
         return cls.__name__
 
 
-def _fill_function(func, globals, defaults, dict, closure_values):
+def _fill_function(func, globals, defaults, dict, module, closure_values):
     """ Fills in the rest of function data into the skeleton function object
         that were created via _make_skel_func().
     """
     func.__globals__.update(globals)
     func.__defaults__ = defaults
     func.__dict__ = dict
+    func.__module__ = module
 
     cells = func.__closure__
     if cells is not None:

tests/cloudpickle_test.py

@@ -697,5 +697,9 @@ def f():
         result = g()
         self.assertEqual(1, result)
 
+    def test_function_module_name(self):
+        func = lambda x: x
+        self.assertEqual(pickle_depickle(func).__module__, func.__module__)
+
 if __name__ == '__main__':
     unittest.main()