You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

455 lines
13 KiB

4 years ago
  1. # encoding: utf-8
  2. """Pickle related utilities. Perhaps this should be called 'can'."""
  3. # Copyright (c) IPython Development Team.
  4. # Distributed under the terms of the Modified BSD License.
  5. import warnings
  6. warnings.warn("ipykernel.pickleutil is deprecated. It has moved to ipyparallel.", DeprecationWarning)
  7. import copy
  8. import sys
  9. from types import FunctionType
  10. try:
  11. import cPickle as pickle
  12. except ImportError:
  13. import pickle
  14. from ipython_genutils import py3compat
  15. from ipython_genutils.importstring import import_item
  16. from ipython_genutils.py3compat import string_types, iteritems, buffer_to_bytes, buffer_to_bytes_py2
  17. # This registers a hook when it's imported
  18. try:
  19. # available since ipyparallel 5.1.1
  20. from ipyparallel.serialize import codeutil
  21. except ImportError:
  22. # Deprecated since ipykernel 4.3.1
  23. from ipykernel import codeutil
  24. from traitlets.log import get_logger
  25. if py3compat.PY3:
  26. buffer = memoryview
  27. class_type = type
  28. else:
  29. from types import ClassType
  30. class_type = (type, ClassType)
  31. try:
  32. PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL
  33. except AttributeError:
  34. PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL
  35. def _get_cell_type(a=None):
  36. """the type of a closure cell doesn't seem to be importable,
  37. so just create one
  38. """
  39. def inner():
  40. return a
  41. return type(py3compat.get_closure(inner)[0])
  42. cell_type = _get_cell_type()
  43. #-------------------------------------------------------------------------------
  44. # Functions
  45. #-------------------------------------------------------------------------------
  46. def interactive(f):
  47. """decorator for making functions appear as interactively defined.
  48. This results in the function being linked to the user_ns as globals()
  49. instead of the module globals().
  50. """
  51. # build new FunctionType, so it can have the right globals
  52. # interactive functions never have closures, that's kind of the point
  53. if isinstance(f, FunctionType):
  54. mainmod = __import__('__main__')
  55. f = FunctionType(f.__code__, mainmod.__dict__,
  56. f.__name__, f.__defaults__,
  57. )
  58. # associate with __main__ for uncanning
  59. f.__module__ = '__main__'
  60. return f
  61. def use_dill():
  62. """use dill to expand serialization support
  63. adds support for object methods and closures to serialization.
  64. """
  65. # import dill causes most of the magic
  66. import dill
  67. # dill doesn't work with cPickle,
  68. # tell the two relevant modules to use plain pickle
  69. global pickle
  70. pickle = dill
  71. try:
  72. from ipykernel import serialize
  73. except ImportError:
  74. pass
  75. else:
  76. serialize.pickle = dill
  77. # disable special function handling, let dill take care of it
  78. can_map.pop(FunctionType, None)
  79. def use_cloudpickle():
  80. """use cloudpickle to expand serialization support
  81. adds support for object methods and closures to serialization.
  82. """
  83. import cloudpickle
  84. global pickle
  85. pickle = cloudpickle
  86. try:
  87. from ipykernel import serialize
  88. except ImportError:
  89. pass
  90. else:
  91. serialize.pickle = cloudpickle
  92. # disable special function handling, let cloudpickle take care of it
  93. can_map.pop(FunctionType, None)
  94. #-------------------------------------------------------------------------------
  95. # Classes
  96. #-------------------------------------------------------------------------------
  97. class CannedObject(object):
  98. def __init__(self, obj, keys=[], hook=None):
  99. """can an object for safe pickling
  100. Parameters
  101. ==========
  102. obj:
  103. The object to be canned
  104. keys: list (optional)
  105. list of attribute names that will be explicitly canned / uncanned
  106. hook: callable (optional)
  107. An optional extra callable,
  108. which can do additional processing of the uncanned object.
  109. large data may be offloaded into the buffers list,
  110. used for zero-copy transfers.
  111. """
  112. self.keys = keys
  113. self.obj = copy.copy(obj)
  114. self.hook = can(hook)
  115. for key in keys:
  116. setattr(self.obj, key, can(getattr(obj, key)))
  117. self.buffers = []
  118. def get_object(self, g=None):
  119. if g is None:
  120. g = {}
  121. obj = self.obj
  122. for key in self.keys:
  123. setattr(obj, key, uncan(getattr(obj, key), g))
  124. if self.hook:
  125. self.hook = uncan(self.hook, g)
  126. self.hook(obj, g)
  127. return self.obj
  128. class Reference(CannedObject):
  129. """object for wrapping a remote reference by name."""
  130. def __init__(self, name):
  131. if not isinstance(name, string_types):
  132. raise TypeError("illegal name: %r"%name)
  133. self.name = name
  134. self.buffers = []
  135. def __repr__(self):
  136. return "<Reference: %r>"%self.name
  137. def get_object(self, g=None):
  138. if g is None:
  139. g = {}
  140. return eval(self.name, g)
  141. class CannedCell(CannedObject):
  142. """Can a closure cell"""
  143. def __init__(self, cell):
  144. self.cell_contents = can(cell.cell_contents)
  145. def get_object(self, g=None):
  146. cell_contents = uncan(self.cell_contents, g)
  147. def inner():
  148. return cell_contents
  149. return py3compat.get_closure(inner)[0]
  150. class CannedFunction(CannedObject):
  151. def __init__(self, f):
  152. self._check_type(f)
  153. self.code = f.__code__
  154. if f.__defaults__:
  155. self.defaults = [ can(fd) for fd in f.__defaults__ ]
  156. else:
  157. self.defaults = None
  158. closure = py3compat.get_closure(f)
  159. if closure:
  160. self.closure = tuple( can(cell) for cell in closure )
  161. else:
  162. self.closure = None
  163. self.module = f.__module__ or '__main__'
  164. self.__name__ = f.__name__
  165. self.buffers = []
  166. def _check_type(self, obj):
  167. assert isinstance(obj, FunctionType), "Not a function type"
  168. def get_object(self, g=None):
  169. # try to load function back into its module:
  170. if not self.module.startswith('__'):
  171. __import__(self.module)
  172. g = sys.modules[self.module].__dict__
  173. if g is None:
  174. g = {}
  175. if self.defaults:
  176. defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
  177. else:
  178. defaults = None
  179. if self.closure:
  180. closure = tuple(uncan(cell, g) for cell in self.closure)
  181. else:
  182. closure = None
  183. newFunc = FunctionType(self.code, g, self.__name__, defaults, closure)
  184. return newFunc
  185. class CannedClass(CannedObject):
  186. def __init__(self, cls):
  187. self._check_type(cls)
  188. self.name = cls.__name__
  189. self.old_style = not isinstance(cls, type)
  190. self._canned_dict = {}
  191. for k,v in cls.__dict__.items():
  192. if k not in ('__weakref__', '__dict__'):
  193. self._canned_dict[k] = can(v)
  194. if self.old_style:
  195. mro = []
  196. else:
  197. mro = cls.mro()
  198. self.parents = [ can(c) for c in mro[1:] ]
  199. self.buffers = []
  200. def _check_type(self, obj):
  201. assert isinstance(obj, class_type), "Not a class type"
  202. def get_object(self, g=None):
  203. parents = tuple(uncan(p, g) for p in self.parents)
  204. return type(self.name, parents, uncan_dict(self._canned_dict, g=g))
  205. class CannedArray(CannedObject):
  206. def __init__(self, obj):
  207. from numpy import ascontiguousarray
  208. self.shape = obj.shape
  209. self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
  210. self.pickled = False
  211. if sum(obj.shape) == 0:
  212. self.pickled = True
  213. elif obj.dtype == 'O':
  214. # can't handle object dtype with buffer approach
  215. self.pickled = True
  216. elif obj.dtype.fields and any(dt == 'O' for dt,sz in obj.dtype.fields.values()):
  217. self.pickled = True
  218. if self.pickled:
  219. # just pickle it
  220. self.buffers = [pickle.dumps(obj, PICKLE_PROTOCOL)]
  221. else:
  222. # ensure contiguous
  223. obj = ascontiguousarray(obj, dtype=None)
  224. self.buffers = [buffer(obj)]
  225. def get_object(self, g=None):
  226. from numpy import frombuffer
  227. data = self.buffers[0]
  228. if self.pickled:
  229. # we just pickled it
  230. return pickle.loads(buffer_to_bytes_py2(data))
  231. else:
  232. if not py3compat.PY3 and isinstance(data, memoryview):
  233. # frombuffer doesn't accept memoryviews on Python 2,
  234. # so cast to old-style buffer
  235. data = buffer(data.tobytes())
  236. return frombuffer(data, dtype=self.dtype).reshape(self.shape)
  237. class CannedBytes(CannedObject):
  238. wrap = staticmethod(buffer_to_bytes)
  239. def __init__(self, obj):
  240. self.buffers = [obj]
  241. def get_object(self, g=None):
  242. data = self.buffers[0]
  243. return self.wrap(data)
  244. class CannedBuffer(CannedBytes):
  245. wrap = buffer
  246. class CannedMemoryView(CannedBytes):
  247. wrap = memoryview
  248. #-------------------------------------------------------------------------------
  249. # Functions
  250. #-------------------------------------------------------------------------------
  251. def _import_mapping(mapping, original=None):
  252. """import any string-keys in a type mapping
  253. """
  254. log = get_logger()
  255. log.debug("Importing canning map")
  256. for key,value in list(mapping.items()):
  257. if isinstance(key, string_types):
  258. try:
  259. cls = import_item(key)
  260. except Exception:
  261. if original and key not in original:
  262. # only message on user-added classes
  263. log.error("canning class not importable: %r", key, exc_info=True)
  264. mapping.pop(key)
  265. else:
  266. mapping[cls] = mapping.pop(key)
  267. def istype(obj, check):
  268. """like isinstance(obj, check), but strict
  269. This won't catch subclasses.
  270. """
  271. if isinstance(check, tuple):
  272. for cls in check:
  273. if type(obj) is cls:
  274. return True
  275. return False
  276. else:
  277. return type(obj) is check
  278. def can(obj):
  279. """prepare an object for pickling"""
  280. import_needed = False
  281. for cls,canner in iteritems(can_map):
  282. if isinstance(cls, string_types):
  283. import_needed = True
  284. break
  285. elif istype(obj, cls):
  286. return canner(obj)
  287. if import_needed:
  288. # perform can_map imports, then try again
  289. # this will usually only happen once
  290. _import_mapping(can_map, _original_can_map)
  291. return can(obj)
  292. return obj
  293. def can_class(obj):
  294. if isinstance(obj, class_type) and obj.__module__ == '__main__':
  295. return CannedClass(obj)
  296. else:
  297. return obj
  298. def can_dict(obj):
  299. """can the *values* of a dict"""
  300. if istype(obj, dict):
  301. newobj = {}
  302. for k, v in iteritems(obj):
  303. newobj[k] = can(v)
  304. return newobj
  305. else:
  306. return obj
  307. sequence_types = (list, tuple, set)
  308. def can_sequence(obj):
  309. """can the elements of a sequence"""
  310. if istype(obj, sequence_types):
  311. t = type(obj)
  312. return t([can(i) for i in obj])
  313. else:
  314. return obj
  315. def uncan(obj, g=None):
  316. """invert canning"""
  317. import_needed = False
  318. for cls,uncanner in iteritems(uncan_map):
  319. if isinstance(cls, string_types):
  320. import_needed = True
  321. break
  322. elif isinstance(obj, cls):
  323. return uncanner(obj, g)
  324. if import_needed:
  325. # perform uncan_map imports, then try again
  326. # this will usually only happen once
  327. _import_mapping(uncan_map, _original_uncan_map)
  328. return uncan(obj, g)
  329. return obj
  330. def uncan_dict(obj, g=None):
  331. if istype(obj, dict):
  332. newobj = {}
  333. for k, v in iteritems(obj):
  334. newobj[k] = uncan(v,g)
  335. return newobj
  336. else:
  337. return obj
  338. def uncan_sequence(obj, g=None):
  339. if istype(obj, sequence_types):
  340. t = type(obj)
  341. return t([uncan(i,g) for i in obj])
  342. else:
  343. return obj
  344. #-------------------------------------------------------------------------------
  345. # API dictionaries
  346. #-------------------------------------------------------------------------------
  347. # These dicts can be extended for custom serialization of new objects
  348. can_map = {
  349. 'numpy.ndarray' : CannedArray,
  350. FunctionType : CannedFunction,
  351. bytes : CannedBytes,
  352. memoryview : CannedMemoryView,
  353. cell_type : CannedCell,
  354. class_type : can_class,
  355. }
  356. if buffer is not memoryview:
  357. can_map[buffer] = CannedBuffer
  358. uncan_map = {
  359. CannedObject : lambda obj, g: obj.get_object(g),
  360. dict : uncan_dict,
  361. }
  362. # for use in _import_mapping:
  363. _original_can_map = can_map.copy()
  364. _original_uncan_map = uncan_map.copy()