Lazy import pickle

numpy · hmaarrfk · Jul 22, 2019 · Jul 22, 2019 · Jul 22, 2019 · Jul 22, 2019
commit 625db1400d332b74fdaa98eb9e5c4c6b193e4293
diff --git a/numpy/compat/py3k.py b/numpy/compat/py3k.py
@@ -8,7 +8,7 @@
            'unicode', 'asunicode', 'asbytes_nested', 'asunicode_nested',
            'asstr', 'open_latin1', 'long', 'basestring', 'sixu',
            'integer_types', 'is_pathlib_path', 'npy_load_module', 'Path',
-           'pickle', 'contextlib_nullcontext', 'os_fspath', 'os_PathLike']
+           'get_pickle', 'contextlib_nullcontext', 'os_fspath', 'os_PathLike']
 
 import sys
 import os
@@ -20,10 +20,12 @@
 if sys.version_info[0] >= 3:
     import io
 
-    try:
-        import pickle5 as pickle
-    except ImportError:
-        import pickle
+    def get_pickle():
+        try:
+            import pickle5 as pickle
+        except ImportError:
+            import pickle
+        return pickle
 
     long = int
     integer_types = (int,)
@@ -58,7 +60,9 @@ def sixu(s):
     strchar = 'U'
 
 else:
-    import cpickle as pickle
+    def get_pickle():
+        import cpickle as pickle
+        return pickle
 
     bytes = str
     long = long

diff --git a/numpy/core/_internal.py b/numpy/core/_internal.py
@@ -909,4 +909,3 @@ def __init__(self, func):
         self.func = func
     def __call__(self, *args, **kwargs):
         return self.func(self, *args, **kwargs)
-
diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
@@ -13,7 +13,7 @@
 from numpy.core import numerictypes as nt
 from numpy.core import _exceptions
 from numpy._globals import _NoValue
-from numpy.compat import pickle, os_fspath, contextlib_nullcontext
+from numpy.compat import os_fspath, contextlib_nullcontext, get_pickle
 
 # save those O(100) nanoseconds!
 umr_maximum = um.maximum.reduce
@@ -233,6 +233,7 @@ def _ptp(a, axis=None, out=None, keepdims=False):
     )
 
 def _dump(self, file, protocol=2):
+    pickle = get_pickle()
     if hasattr(file, 'write'):
         ctx = contextlib_nullcontext(file)
     else:
@@ -241,4 +242,5 @@ def _dump(self, file, protocol=2):
         pickle.dump(self, f, protocol=protocol)
 
 def _dumps(self, protocol=2):
+    pickle = get_pickle()
     return pickle.dumps(self, protocol=protocol)
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
@@ -9,7 +9,7 @@
 import contextlib
 
 import numpy as np
-from numpy.compat import pickle, basestring
+from numpy.compat import get_pickle, basestring
 from . import multiarray
 from .multiarray import (
     _fastCopyAndTranspose as fastCopyAndTranspose, ALLOW_THREADS,
@@ -49,6 +49,7 @@
 
 
 def loads(*args, **kwargs):
+    pickle = get_pickle()
     # NumPy 1.15.0, 2017-12-10
     warnings.warn(
         "np.core.numeric.loads is deprecated, use pickle.loads instead",
@@ -937,7 +938,7 @@ def tensordot(a, b, axes=2):
     Returns
     -------
     output : ndarray
-        The tensor dot product of the input.  
+        The tensor dot product of the input.
 
     See Also
     --------
@@ -2038,6 +2039,7 @@ def load(file):
     load, save
 
     """
+    pickle = get_pickle()
     # NumPy 1.15.0, 2017-12-10
     warnings.warn(
         "np.core.numeric.load is deprecated, use pickle.load instead",

diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
@@ -9,7 +9,8 @@
     assert_, assert_equal, assert_raises, assert_warns, suppress_warnings,
     assert_raises_regex,
     )
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 # Use pytz to test out various time zones if available
 try:

diff --git a/numpy/core/tests/test_dtype.py b/numpy/core/tests/test_dtype.py
@@ -10,7 +10,8 @@
 from numpy.core._rational_tests import rational
 from numpy.testing import (
     assert_, assert_equal, assert_array_equal, assert_raises, HAS_REFCOUNT)
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 from itertools import permutations
 
 def assert_dtype_equal(a, b):
@@ -138,11 +139,11 @@ def test_bad_param(self):
                          'offsets':[0, 2]}, align=True)
 
     def test_field_order_equality(self):
-        x = np.dtype({'names': ['A', 'B'], 
-                      'formats': ['i4', 'f4'], 
+        x = np.dtype({'names': ['A', 'B'],
+                      'formats': ['i4', 'f4'],
                       'offsets': [0, 4]})
-        y = np.dtype({'names': ['B', 'A'], 
-                      'formats': ['f4', 'i4'], 
+        y = np.dtype({'names': ['B', 'A'],
+                      'formats': ['f4', 'i4'],
                       'offsets': [4, 0]})
         assert_equal(x == y, False)
 
@@ -1279,4 +1280,3 @@ def test_pairs(self, pair):
         pair_type = np.dtype('{},{}'.format(*pair))
         expected = np.dtype([('f0', pair[0]), ('f1', pair[1])])
         assert_equal(pair_type, expected)
-
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -21,7 +21,8 @@
 import pytest
 from contextlib import contextmanager
 
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 try:
     import pathlib
@@ -114,7 +115,7 @@ def test_writeable_any_base(self):
         # Ensure that any base being writeable is sufficient to change flag;
         # this is especially interesting for arrays from an array interface.
         arr = np.arange(10)
-        
+
         class subclass(np.ndarray):
             pass
 
@@ -3967,13 +3968,13 @@ def test_subarray_int_shape(self):
 
     def test_datetime64_byteorder(self):
         original = np.array([['2015-02-24T00:00:00.000000000']], dtype='datetime64[ns]')
-    
+
         original_byte_reversed = original.copy(order='K')
         original_byte_reversed.dtype = original_byte_reversed.dtype.newbyteorder('S')
         original_byte_reversed.byteswap(inplace=True)
 
         new = pickle.loads(pickle.dumps(original_byte_reversed))
-    
+
         assert_equal(original.dtype, new.dtype)
 
 
@@ -4868,7 +4869,7 @@ def test_fromfile_offset(self):
             offset_bytes = self.dtype.itemsize
             z = np.fromfile(f, dtype=self.dtype, offset=offset_bytes)
             assert_array_equal(z, self.x.flat[offset_items+count_items+1:])
-        
+
         with open(self.filename, 'wb') as f:
             self.x.tofile(f, sep=",")
 
@@ -6225,14 +6226,14 @@ def test_dot_equivalent(self, args):
 
         r3 = np.matmul(args[0].copy(), args[1].copy())
         assert_equal(r1, r3)
-    
+
     def test_matmul_object(self):
         import fractions
 
         f = np.vectorize(fractions.Fraction)
         def random_ints():
             return np.random.randint(1, 1000, size=(10, 3, 3))
-        M1 = f(random_ints(), random_ints()) 
+        M1 = f(random_ints(), random_ints())
         M2 = f(random_ints(), random_ints())
 
         M3 = self.matmul(M1, M2)

diff --git a/numpy/core/tests/test_overrides.py b/numpy/core/tests/test_overrides.py
@@ -10,7 +10,8 @@
 from numpy.core.overrides import (
     _get_implementing_args, array_function_dispatch,
     verify_matching_signatures, ARRAY_FUNCTION_ENABLED)
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 import pytest
 
 

diff --git a/numpy/core/tests/test_records.py b/numpy/core/tests/test_records.py
@@ -17,8 +17,8 @@
     assert_, assert_equal, assert_array_equal, assert_array_almost_equal,
     assert_raises, temppath
     )
-from numpy.compat import pickle
-
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 class TestFromrecords(object):
     def test_fromrecords(self):

diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
@@ -16,7 +16,8 @@
         assert_raises_regex, assert_warns, suppress_warnings,
         _assert_valid_refcount, HAS_REFCOUNT,
         )
-from numpy.compat import asbytes, asunicode, long, pickle
+from numpy.compat import asbytes, asunicode, long, get_pickle
+pickle = get_pickle()
 
 try:
     RecursionError

diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
@@ -15,8 +15,8 @@
     assert_almost_equal, assert_array_almost_equal, assert_no_warnings,
     assert_allclose,
     )
-from numpy.compat import pickle
-
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 class TestUfuncKwargs(object):
     def test_kwarg_exact(self):

diff --git a/numpy/lib/format.py b/numpy/lib/format.py
@@ -169,7 +169,7 @@
 import warnings
 from numpy.lib.utils import safe_eval
 from numpy.compat import (
-    isfileobj, long, os_fspath, pickle
+    isfileobj, long, os_fspath, get_pickle
     )
 
 
@@ -656,6 +656,7 @@ def write_array(fp, array, version=None, allow_pickle=True, pickle_kwargs=None):
                              "allow_pickle=False")
         if pickle_kwargs is None:
             pickle_kwargs = {}
+        pickle = get_pickle()
         pickle.dump(array, fp, protocol=3, **pickle_kwargs)
     elif array.flags.f_contiguous and not array.flags.c_contiguous:
         if isfileobj(fp):
@@ -724,6 +725,7 @@ def read_array(fp, allow_pickle=False, pickle_kwargs=None):
         if pickle_kwargs is None:
             pickle_kwargs = {}
         try:
+            pickle = get_pickle()
             array = pickle.load(fp, **pickle_kwargs)
         except UnicodeError as err:
             if sys.version_info[0] >= 3:

diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py
@@ -25,7 +25,7 @@
 
 from numpy.compat import (
     asbytes, asstr, asunicode, bytes, basestring, os_fspath, os_PathLike,
-    pickle, contextlib_nullcontext
+    get_pickle, contextlib_nullcontext
     )
 
 if sys.version_info[0] >= 3:
@@ -37,6 +37,7 @@
 
 @set_module('numpy')
 def loads(*args, **kwargs):
+    pickle = get_pickle()
     # NumPy 1.15.0, 2017-12-10
     warnings.warn(
         "np.loads is deprecated, use pickle.loads instead",
@@ -457,6 +458,7 @@ def load(file, mmap_mode=None, allow_pickle=False, fix_imports=True,
                 raise ValueError("Cannot load file containing pickled data "
                                  "when allow_pickle=False")
             try:
+                pickle = get_pickle()
                 return pickle.load(fid, **pickle_kwargs)
             except Exception:
                 raise IOError(
@@ -680,7 +682,7 @@ def savez_compressed(file, *args, **kwds):
     The ``.npz`` file format is a zipped archive of files named after the
     variables they contain.  The archive is compressed with
     ``zipfile.ZIP_DEFLATED`` and each file in the archive contains one variable
-    in ``.npy`` format. For a description of the ``.npy`` format, see 
+    in ``.npy`` format. For a description of the ``.npy`` format, see
     :py:mod:`numpy.lib.format`.
 
 

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -46,7 +46,7 @@
 from numpy import expand_dims
 from numpy.core.numeric import normalize_axis_tuple
 from numpy.core._internal import recursive
-from numpy.compat import pickle
+from numpy.compat import get_pickle
 
 
 __all__ = [
@@ -7919,6 +7919,7 @@ def dump(a, F):
 
     """
     _pickle_warn('dump')
+    pickle = get_pickle()
     if not hasattr(F, 'readline'):
         with open(F, 'w') as F:
             pickle.dump(a, F)
@@ -7940,6 +7941,7 @@ def dumps(a):
 
     """
     _pickle_warn('dumps')
+    pickle = get_pickle()
     return pickle.dumps(a)
 
 

diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
@@ -49,8 +49,8 @@
     ravel, repeat, reshape, resize, shape, sin, sinh, sometrue, sort, sqrt,
     subtract, sum, take, tan, tanh, transpose, where, zeros,
     )
-from numpy.compat import pickle
-
+from numpy.compat import get_pickle
+pickle = get_pickle()
 pi = np.pi
 
 

diff --git a/numpy/ma/tests/test_mrecords.py b/numpy/ma/tests/test_mrecords.py
@@ -23,7 +23,8 @@
     assert_, assert_equal,
     assert_equal_records,
     )
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 
 class TestMRecords(object):

diff --git a/numpy/ma/tests/test_old_ma.py b/numpy/ma/tests/test_old_ma.py
@@ -21,7 +21,8 @@
     repeat, resize, shape, sin, sinh, sometrue, sort, sqrt, subtract, sum,
     take, tan, tanh, transpose, where, zeros,
     )
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 pi = np.pi
 

diff --git a/numpy/matrixlib/tests/test_masked_matrix.py b/numpy/matrixlib/tests/test_masked_matrix.py
@@ -7,7 +7,8 @@
                            MaskType, getmask, MaskedArray, nomask,
                            log, add, hypot, divide)
 from numpy.ma.extras import mr_
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
 
 
 class MMatrix(MaskedArray, np.matrix,):

diff --git a/numpy/tests/test_reloading.py b/numpy/tests/test_reloading.py
@@ -3,7 +3,9 @@
 import sys
 
 from numpy.testing import assert_raises, assert_, assert_equal
-from numpy.compat import pickle
+from numpy.compat import get_pickle
+pickle = get_pickle()
+
 
 if sys.version_info[:2] >= (3, 4):
     from importlib import reload
Original file line number	Diff line number	Diff line change
Expand Up		@@ -909,4 +909,3 @@ def __init__(self, func):
		self.func = func
		def __call__(self, args, *kwargs):
		return self.func(self, args, *kwargs)