github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/internal/pickler.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Pickler for values, functions, and classes.
    19  
    20  For internal use only. No backwards compatibility guarantees.
    21  
    22  Pickles created by the pickling library contain non-ASCII characters, so
    23  we base64-encode the results so that we can put them in a JSON objects.
    24  The pickler is used to embed FlatMap callable objects into the workflow JSON
    25  description.
    26  
    27  The pickler module should be used to pickle functions and modules; for values,
    28  the coders.*PickleCoder classes should be used instead.
    29  """
    30  
    31  from apache_beam.internal import cloudpickle_pickler
    32  from apache_beam.internal import dill_pickler
    33  
    34  USE_CLOUDPICKLE = 'cloudpickle'
    35  USE_DILL = 'dill'
    36  DEFAULT_PICKLE_LIB = USE_DILL
    37  
    38  desired_pickle_lib = dill_pickler
    39  
    40  
    41  def dumps(o, enable_trace=True, use_zlib=False):
    42    # type: (...) -> bytes
    43  
    44    return desired_pickle_lib.dumps(
    45        o, enable_trace=enable_trace, use_zlib=use_zlib)
    46  
    47  
    48  def loads(encoded, enable_trace=True, use_zlib=False):
    49    """For internal use only; no backwards-compatibility guarantees."""
    50  
    51    return desired_pickle_lib.loads(
    52        encoded, enable_trace=enable_trace, use_zlib=use_zlib)
    53  
    54  
    55  def dump_session(file_path):
    56    """For internal use only; no backwards-compatibility guarantees.
    57  
    58    Pickle the current python session to be used in the worker.
    59    """
    60  
    61    return desired_pickle_lib.dump_session(file_path)
    62  
    63  
    64  def load_session(file_path):
    65    return desired_pickle_lib.load_session(file_path)
    66  
    67  
    68  def set_library(selected_library=DEFAULT_PICKLE_LIB):
    69    """ Sets pickle library that will be used. """
    70    global desired_pickle_lib
    71    # If switching to or from dill, update the pickler hook overrides.
    72    if (selected_library == USE_DILL) != (desired_pickle_lib == dill_pickler):
    73      dill_pickler.override_pickler_hooks(selected_library == USE_DILL)
    74  
    75    if selected_library == 'default':
    76      selected_library = DEFAULT_PICKLE_LIB
    77  
    78    if selected_library == USE_DILL:
    79      desired_pickle_lib = dill_pickler
    80    elif selected_library == USE_CLOUDPICKLE:
    81      desired_pickle_lib = cloudpickle_pickler
    82    else:
    83      raise ValueError(f'Unknown pickler library: {selected_library}')