1  """Pickle related utilities. Perhaps this should be called 'can'."""
   2  
   3  # Copyright (c) IPython Development Team.
   4  # Distributed under the terms of the Modified BSD License.
   5  import copy
   6  import pickle
   7  import sys
   8  import typing
   9  import warnings
  10  from types import FunctionType
  11  
  12  # This registers a hook when it's imported
  13  try:
  14      from ipyparallel.serialize import codeutil  # noqa: F401
  15  except ImportError:
  16      pass
  17  from traitlets.log import get_logger
  18  from traitlets.utils.importstring import import_item
  19  
  20  warnings.warn(
  21      "ipykernel.pickleutil is deprecated. It has moved to ipyparallel.",
  22      DeprecationWarning,
  23      stacklevel=2,
  24  )
  25  
  26  buffer = memoryview
  27  class_type = type
  28  
  29  PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL
  30  
  31  
  32  def _get_cell_type(a=None):
  33      """the type of a closure cell doesn't seem to be importable,
  34      so just create one
  35      """
  36  
  37      def inner():
  38          return a
  39  
  40      return type(inner.__closure__[0])  # type:ignore[index]
  41  
  42  
  43  cell_type = _get_cell_type()
  44  
  45  # -------------------------------------------------------------------------------
  46  # Functions
  47  # -------------------------------------------------------------------------------
  48  
  49  
  50  def interactive(f):
  51      """decorator for making functions appear as interactively defined.
  52      This results in the function being linked to the user_ns as globals()
  53      instead of the module globals().
  54      """
  55  
  56      # build new FunctionType, so it can have the right globals
  57      # interactive functions never have closures, that's kind of the point
  58      if isinstance(f, FunctionType):
  59          mainmod = __import__("__main__")
  60          f = FunctionType(
  61              f.__code__,
  62              mainmod.__dict__,
  63              f.__name__,
  64              f.__defaults__,
  65          )
  66      # associate with __main__ for uncanning
  67      f.__module__ = "__main__"
  68      return f
  69  
  70  
  71  def use_dill():
  72      """use dill to expand serialization support
  73  
  74      adds support for object methods and closures to serialization.
  75      """
  76      # import dill causes most of the magic
  77      import dill
  78  
  79      # dill doesn't work with cPickle,
  80      # tell the two relevant modules to use plain pickle
  81  
  82      global pickle  # noqa: PLW0603
  83      pickle = dill
  84  
  85      try:
  86          from ipykernel import serialize
  87      except ImportError:
  88          pass
  89      else:
  90          serialize.pickle = dill  # type:ignore[attr-defined]
  91  
  92      # disable special function handling, let dill take care of it
  93      can_map.pop(FunctionType, None)
  94  
  95  
  96  def use_cloudpickle():
  97      """use cloudpickle to expand serialization support
  98  
  99      adds support for object methods and closures to serialization.
 100      """
 101      import cloudpickle
 102  
 103      global pickle  # noqa: PLW0603
 104      pickle = cloudpickle
 105  
 106      try:
 107          from ipykernel import serialize
 108      except ImportError:
 109          pass
 110      else:
 111          serialize.pickle = cloudpickle  # type:ignore[attr-defined]
 112  
 113      # disable special function handling, let cloudpickle take care of it
 114      can_map.pop(FunctionType, None)
 115  
 116  
 117  # -------------------------------------------------------------------------------
 118  # Classes
 119  # -------------------------------------------------------------------------------
 120  
 121  
 122  class CannedObject:
 123      """A canned object."""
 124  
 125      def __init__(self, obj, keys=None, hook=None):
 126          """can an object for safe pickling
 127  
 128          Parameters
 129          ----------
 130          obj
 131              The object to be canned
 132          keys : list (optional)
 133              list of attribute names that will be explicitly canned / uncanned
 134          hook : callable (optional)
 135              An optional extra callable,
 136              which can do additional processing of the uncanned object.
 137  
 138          Notes
 139          -----
 140          large data may be offloaded into the buffers list,
 141          used for zero-copy transfers.
 142          """
 143          self.keys = keys or []
 144          self.obj = copy.copy(obj)
 145          self.hook = can(hook)
 146          for key in keys:
 147              setattr(self.obj, key, can(getattr(obj, key)))
 148  
 149          self.buffers = []
 150  
 151      def get_object(self, g=None):
 152          """Get an object."""
 153          if g is None:
 154              g = {}
 155          obj = self.obj
 156          for key in self.keys:
 157              setattr(obj, key, uncan(getattr(obj, key), g))
 158  
 159          if self.hook:
 160              self.hook = uncan(self.hook, g)
 161              self.hook(obj, g)
 162          return self.obj
 163  
 164  
 165  class Reference(CannedObject):
 166      """object for wrapping a remote reference by name."""
 167  
 168      def __init__(self, name):
 169          """Initialize the reference."""
 170          if not isinstance(name, str):
 171              raise TypeError("illegal name: %r" % name)
 172          self.name = name
 173          self.buffers = []
 174  
 175      def __repr__(self):
 176          """Get the string repr of the reference."""
 177          return "<Reference: %r>" % self.name
 178  
 179      def get_object(self, g=None):
 180          """Get an object in the reference."""
 181          if g is None:
 182              g = {}
 183  
 184          return eval(self.name, g)
 185  
 186  
 187  class CannedCell(CannedObject):
 188      """Can a closure cell"""
 189  
 190      def __init__(self, cell):
 191          """Initialize the canned cell."""
 192          self.cell_contents = can(cell.cell_contents)
 193  
 194      def get_object(self, g=None):
 195          """Get an object in the cell."""
 196          cell_contents = uncan(self.cell_contents, g)
 197  
 198          def inner():
 199              """Inner function."""
 200              return cell_contents
 201  
 202          return inner.__closure__[0]  # type:ignore[index]
 203  
 204  
 205  class CannedFunction(CannedObject):
 206      """Can a function."""
 207  
 208      def __init__(self, f):
 209          """Initialize the can"""
 210          self._check_type(f)
 211          self.code = f.__code__
 212          self.defaults: typing.Optional[typing.List[typing.Any]]
 213          if f.__defaults__:
 214              self.defaults = [can(fd) for fd in f.__defaults__]
 215          else:
 216              self.defaults = None
 217  
 218          self.closure: typing.Any
 219          closure = f.__closure__
 220          if closure:
 221              self.closure = tuple(can(cell) for cell in closure)
 222          else:
 223              self.closure = None
 224  
 225          self.module = f.__module__ or "__main__"
 226          self.__name__ = f.__name__
 227          self.buffers = []
 228  
 229      def _check_type(self, obj):
 230          assert isinstance(obj, FunctionType), "Not a function type"
 231  
 232      def get_object(self, g=None):
 233          """Get an object out of the can."""
 234          # try to load function back into its module:
 235          if not self.module.startswith("__"):
 236              __import__(self.module)
 237              g = sys.modules[self.module].__dict__
 238  
 239          if g is None:
 240              g = {}
 241          defaults = tuple(uncan(cfd, g) for cfd in self.defaults) if self.defaults else None
 242          closure = tuple(uncan(cell, g) for cell in self.closure) if self.closure else None
 243          return FunctionType(self.code, g, self.__name__, defaults, closure)
 244  
 245  
 246  class CannedClass(CannedObject):
 247      """A canned class object."""
 248  
 249      def __init__(self, cls):
 250          """Initialize the can."""
 251          self._check_type(cls)
 252          self.name = cls.__name__
 253          self.old_style = not isinstance(cls, type)
 254          self._canned_dict = {}
 255          for k, v in cls.__dict__.items():
 256              if k not in ("__weakref__", "__dict__"):
 257                  self._canned_dict[k] = can(v)
 258          mro = [] if self.old_style else cls.mro()
 259  
 260          self.parents = [can(c) for c in mro[1:]]
 261          self.buffers = []
 262  
 263      def _check_type(self, obj):
 264          assert isinstance(obj, class_type), "Not a class type"
 265  
 266      def get_object(self, g=None):
 267          """Get an object from the can."""
 268          parents = tuple(uncan(p, g) for p in self.parents)
 269          return type(self.name, parents, uncan_dict(self._canned_dict, g=g))
 270  
 271  
 272  class CannedArray(CannedObject):
 273      """A canned numpy array."""
 274  
 275      def __init__(self, obj):
 276          """Initialize the can."""
 277          from numpy import ascontiguousarray
 278  
 279          self.shape = obj.shape
 280          self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
 281          self.pickled = False
 282          if sum(obj.shape) == 0:
 283              self.pickled = True
 284          elif obj.dtype == "O":
 285              # can't handle object dtype with buffer approach
 286              self.pickled = True
 287          elif obj.dtype.fields and any(dt == "O" for dt, sz in obj.dtype.fields.values()):
 288              self.pickled = True
 289          if self.pickled:
 290              # just pickle it
 291              self.buffers = [pickle.dumps(obj, PICKLE_PROTOCOL)]
 292          else:
 293              # ensure contiguous
 294              obj = ascontiguousarray(obj, dtype=None)
 295              self.buffers = [buffer(obj)]
 296  
 297      def get_object(self, g=None):
 298          """Get the object."""
 299          from numpy import frombuffer
 300  
 301          data = self.buffers[0]
 302          if self.pickled:
 303              # we just pickled it
 304              return pickle.loads(data)
 305          return frombuffer(data, dtype=self.dtype).reshape(self.shape)
 306  
 307  
 308  class CannedBytes(CannedObject):
 309      """A canned bytes object."""
 310  
 311      @staticmethod
 312      def wrap(buf: typing.Union[memoryview, bytes, typing.SupportsBytes]) -> bytes:
 313          """Cast a buffer or memoryview object to bytes"""
 314          if isinstance(buf, memoryview):
 315              return buf.tobytes()
 316          if not isinstance(buf, bytes):
 317              return bytes(buf)
 318          return buf
 319  
 320      def __init__(self, obj):
 321          """Initialize the can."""
 322          self.buffers = [obj]
 323  
 324      def get_object(self, g=None):
 325          """Get the canned object."""
 326          data = self.buffers[0]
 327          return self.wrap(data)
 328  
 329  
 330  class CannedBuffer(CannedBytes):
 331      """A canned buffer."""
 332  
 333      wrap = buffer  # type:ignore[assignment]
 334  
 335  
 336  class CannedMemoryView(CannedBytes):
 337      """A canned memory view."""
 338  
 339      wrap = memoryview  # type:ignore[assignment]
 340  
 341  
 342  # -------------------------------------------------------------------------------
 343  # Functions
 344  # -------------------------------------------------------------------------------
 345  
 346  
 347  def _import_mapping(mapping, original=None):
 348      """import any string-keys in a type mapping"""
 349      log = get_logger()
 350      log.debug("Importing canning map")
 351      for key, _ in list(mapping.items()):
 352          if isinstance(key, str):
 353              try:
 354                  cls = import_item(key)
 355              except Exception:
 356                  if original and key not in original:
 357                      # only message on user-added classes
 358                      log.error("canning class not importable: %r", key, exc_info=True)  # noqa: G201
 359                  mapping.pop(key)
 360              else:
 361                  mapping[cls] = mapping.pop(key)
 362  
 363  
 364  def istype(obj, check):
 365      """like isinstance(obj, check), but strict
 366  
 367      This won't catch subclasses.
 368      """
 369      if isinstance(check, tuple):
 370          return any(type(obj) is cls for cls in check)
 371      return type(obj) is check
 372  
 373  
 374  def can(obj):
 375      """prepare an object for pickling"""
 376  
 377      import_needed = False
 378  
 379      for cls, canner in can_map.items():
 380          if isinstance(cls, str):
 381              import_needed = True
 382              break
 383          if istype(obj, cls):
 384              return canner(obj)
 385  
 386      if import_needed:
 387          # perform can_map imports, then try again
 388          # this will usually only happen once
 389          _import_mapping(can_map, _original_can_map)
 390          return can(obj)
 391  
 392      return obj
 393  
 394  
 395  def can_class(obj):
 396      """Can a class object."""
 397      if isinstance(obj, class_type) and obj.__module__ == "__main__":
 398          return CannedClass(obj)
 399      return obj
 400  
 401  
 402  def can_dict(obj):
 403      """can the *values* of a dict"""
 404      if istype(obj, dict):
 405          newobj = {}
 406          for k, v in obj.items():
 407              newobj[k] = can(v)
 408          return newobj
 409      return obj
 410  
 411  
 412  sequence_types = (list, tuple, set)
 413  
 414  
 415  def can_sequence(obj):
 416      """can the elements of a sequence"""
 417      if istype(obj, sequence_types):
 418          t = type(obj)
 419          return t([can(i) for i in obj])
 420      return obj
 421  
 422  
 423  def uncan(obj, g=None):
 424      """invert canning"""
 425  
 426      import_needed = False
 427      for cls, uncanner in uncan_map.items():
 428          if isinstance(cls, str):
 429              import_needed = True
 430              break
 431          if isinstance(obj, cls):
 432              return uncanner(obj, g)
 433  
 434      if import_needed:
 435          # perform uncan_map imports, then try again
 436          # this will usually only happen once
 437          _import_mapping(uncan_map, _original_uncan_map)
 438          return uncan(obj, g)
 439  
 440      return obj
 441  
 442  
 443  def uncan_dict(obj, g=None):
 444      """Uncan a dict object."""
 445      if istype(obj, dict):
 446          newobj = {}
 447          for k, v in obj.items():
 448              newobj[k] = uncan(v, g)
 449          return newobj
 450      return obj
 451  
 452  
 453  def uncan_sequence(obj, g=None):
 454      """Uncan a sequence."""
 455      if istype(obj, sequence_types):
 456          t = type(obj)
 457          return t([uncan(i, g) for i in obj])
 458      return obj
 459  
 460  
 461  # -------------------------------------------------------------------------------
 462  # API dictionaries
 463  # -------------------------------------------------------------------------------
 464  
 465  # These dicts can be extended for custom serialization of new objects
 466  
 467  can_map = {
 468      "numpy.ndarray": CannedArray,
 469      FunctionType: CannedFunction,
 470      bytes: CannedBytes,
 471      memoryview: CannedMemoryView,
 472      cell_type: CannedCell,
 473      class_type: can_class,
 474  }
 475  if buffer is not memoryview:
 476      can_map[buffer] = CannedBuffer
 477  
 478  uncan_map: typing.Dict[type, typing.Any] = {
 479      CannedObject: lambda obj, g: obj.get_object(g),
 480      dict: uncan_dict,
 481  }
 482  
 483  # for use in _import_mapping:
 484  _original_can_map = can_map.copy()
 485  _original_uncan_map = uncan_map.copy()
 486