用pickle或dill序列化__main__中的對象

Question

我有酸洗問題。 我想在我的主腳本中序列化一個函數，然后加載它並在另一個腳本中運行它。 為了證明這一點，我制作了 2 個腳本：

嘗試 1：天真的方法：

dill_pickle_script_1.py

import pickle
import time

def my_func(a, b):
    time.sleep(0.1)  # The purpose of this will become evident at the end
    return a+b

if __name__ == '__main__':
    with open('testfile.pkl', 'wb') as f:
        pickle.dump(my_func, f)

dill_pickle_script_2.py

import pickle

if __name__ == '__main__':
    with open('testfile.pkl') as f:
        func = pickle.load(f)
        assert func(1, 2)==3

問題：當我運行腳本 2 時，我得到AttributeError: 'module' object has no attribute 'my_func' 。 我明白為什么：因為當 my_func 在 script1 中被序列化時，它屬於__main__模塊。 dill_pickle_script_2 不知道那里的__main__引用了 dill_pickle_script_1 的命名空間，因此找不到引用。

嘗試 2：插入絕對導入

我通過添加一個小技巧來解決這個問題 - 在腌制之前，我在 dill_pickle_script_1 中添加了一個對 my_func 的絕對導入。

dill_pickle_script_1.py

import pickle
import time

def my_func(a, b):
    time.sleep(0.1)
    return a+b

if __name__ == '__main__':
    from dill_pickle_script_1 import my_func  # Added absolute import
    with open('testfile.pkl', 'wb') as f:
        pickle.dump(my_func, f)

現在它起作用了！ 但是，我想避免每次想要這樣做時都必須這樣做。 （另外，我想讓我的酸洗在其他一些不知道 my_func 來自哪個模塊的模塊中完成）。

嘗試 3：蒔蘿

我認為包dill允許您在 main 中序列化內容並將它們加載到其他地方。 所以我試過了：

dill_pickle_script_1.py

import dill
import time

def my_func(a, b):
    time.sleep(0.1)
    return a+b

if __name__ == '__main__':
    with open('testfile.pkl', 'wb') as f:
        dill.dump(my_func, f)

dill_pickle_script_2.py

import dill

if __name__ == '__main__':
    with open('testfile.pkl') as f:
        func = dill.load(f)
        assert func(1, 2)==3

但是，現在我遇到了另一個問題：運行dill_pickle_script_2.py ，出現NameError: global name 'time' is not defined 。 似乎 dill 沒有意識到 my_func 引用了time模塊並且必須在加載時導入它。

我的問題？

如何在 main 中序列化一個對象，並在另一個腳本中再次加載它，以便該對象使用的所有導入也被加載，而無需在嘗試 2 中進行令人討厭的小技巧？

Answer 1

嗯，我找到了解決辦法。 這是一個可怕但整潔的雜物，並不能保證在所有情況下都有效。 歡迎提出任何改進建議。 解決方案包括使用 pickle 字符串中的絕對模塊引用替換主引用，使用以下幫助函數：

import sys
import os

def pickle_dumps_without_main_refs(obj):
    """
    Yeah this is horrible, but it allows you to pickle an object in the main module so that it can be reloaded in another
    module.
    :param obj:
    :return:
    """
    currently_run_file = sys.argv[0]
    module_path = file_path_to_absolute_module(currently_run_file)
    pickle_str = pickle.dumps(obj, protocol=0)
    pickle_str = pickle_str.replace('__main__', module_path)  # Hack!
    return pickle_str


def pickle_dump_without_main_refs(obj, file_obj):
    string = pickle_dumps_without_main_refs(obj)
    file_obj.write(string)


def file_path_to_absolute_module(file_path):
    """
    Given a file path, return an import path.
    :param file_path: A file path.
    :return:
    """
    assert os.path.exists(file_path)
    file_loc, ext = os.path.splitext(file_path)
    assert ext in ('.py', '.pyc')
    directory, module = os.path.split(file_loc)
    module_path = [module]
    while True:
        if os.path.exists(os.path.join(directory, '__init__.py')):
            directory, package = os.path.split(directory)
            module_path.append(package)
        else:
            break
    path = '.'.join(module_path[::-1])
    return path

現在，我可以簡單地更改dill_pickle_script_1.py說

import time
from artemis.remote.child_processes import pickle_dump_without_main_refs


def my_func(a, b):
    time.sleep(0.1)
    return a+b

if __name__ == '__main__':
    with open('testfile.pkl', 'wb') as f:
        pickle_dump_without_main_refs(my_func, f)

然后dill_pickle_script_2.py工作！

Answer 2

您可以將dill.dump與recurse=True或dill.settings["recurse"] = True 。 它將捕獲閉包：

在文件 A 中：

import time
import dill

def my_func(a, b):
  time.sleep(0.1)
  return a + b

with open("tmp.pkl", "wb") as f:
  dill.dump(my_func, f, recurse=True)

在文件 B 中：

import dill

with open("tmp.pkl", "rb") as f:
  my_func = dill.load(f)