Serializing custom modules together with object in python

Question

Problem

Let's say I have this module named custom_module:

class CustomClass:
    pass

And I use this class in a script that serializes an object of the class I defined in custom_module:

import cloudpickle
import custom_module as cm

custom_object = cm.CustomClass()
with open('filename.pkl', 'wb') as file:
    cloudpickle.dump(custom_object, file)

I copy this pickle file to another environment and load it with another script:

import cloudpickle

with open('filename.pkl', 'rb') as file:
    custom_object = cloudpickle.load(file)

That would yield this error:

Traceback (most recent call last):
  File "loader.py", line 4, in <module>
    custom_object = cloudpickle.load(file)
ModuleNotFoundError: No module named 'custom_module'

---

Awkward solution

As a workaround it is possible to read and execute everything that is in custom_module in my script:

exec(open(path.join('custom_module.py')).read())

but that looks really weird and I can't use CustomClass as cm.CustomClass. Is there any other solution that doesn't involve copying all the code in the first environment to the second?

Answer 1

You can solve the problem by reimplement CustomClass in custom_module the following way:

def __CustomClass():
    class CustomeClass:
        ... # Your implementaion here
    return CustomeClass

CustomClass = __CustomClass()

The error will gone if you are lucky enough. If not, you need to dive into CustomClass to find out other functions or classes that are defined in other local modules, and reimplement them using the same method.

You can find more detail in this question. You may also use cloudpickle.register_pickle_by_value to solve the issue, but it is labeled as an experimental feature.