我有一个 pickle 问题。我想在我的主脚本中序列化一个函数,然后加载它并在另一个脚本中运行它。为了证明这一点,我制作了 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'
。我明白为什么了:因为在script1中序列化my_func时,它属于__main__
模块。 dill_pickle_script_2 不知道 __main__
那里引用了 dill_pickle_script_1 的命名空间,因此找不到引用。
尝试 2:插入绝对导入
我通过添加一些 hack 来解决这个问题 - 我在 pickle 之前向 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)
现在可以了!但是,我想避免每次我想这样做时都必须这样做。 (另外,我想让我的 pickle 在其他一些模块中完成,这些模块不知道 my_func 来自哪个模块)。
尝试 3: dill
我头那个包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 中进行令人讨厌的小 hack?
最佳答案
好吧,我找到了解决办法。这是一个可怕但整洁的拼凑,并不能保证在所有情况下都能正常工作。欢迎提出任何改进建议。解决方案包括使用以下辅助函数将 main 引用替换为 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
改成 say
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
起作用了!
关于python - 使用 pickle 或 dill 序列化 __main__ 中的对象,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/45616584/