MAINT, TST import pickle from numpy.core.numeric

All imports of pickle from numpy modules are now done this way:
>>> from numpy.core.numeric import pickle

Also, some loops on protocol numbers are added over pickle tests that
were not caught from numpy#12090

doc/cdoc/numpyfilter.py

@@ -14,10 +14,7 @@
 import textwrap
 import optparse
 
-if sys.version_info[0] >= 3:
-    import pickle
-else:
-    import cPickle as pickle
+from numpy.core.numeric import pickle
 
 CACHE_FILE = 'build/rst-cache.pck'
 

numpy/core/numeric.py

@@ -40,7 +40,13 @@
 newaxis = None
 
 if sys.version_info[0] >= 3:
-    import pickle
+    if sys.version_info[1] in (6, 7):
+        try:
+            import pickle5 as pickle
+        except ImportError:
+            import pickle
+    else:
+        import pickle
     basestring = str
     import builtins
 else:

numpy/core/tests/test_datetime.py

@@ -1,6 +1,5 @@
 from __future__ import division, absolute_import, print_function
 
-import pickle
 
 import numpy
 import numpy as np
@@ -9,6 +8,7 @@
 from numpy.testing import (
     assert_, assert_equal, assert_raises, assert_warns, suppress_warnings,
     )
+from numpy.core.numeric import pickle
 
 # Use pytz to test out various time zones if available
 try:

numpy/core/tests/test_dtype.py

@@ -1,6 +1,5 @@
 from __future__ import division, absolute_import, print_function
 
-import pickle
 import sys
 import operator
 import pytest
@@ -9,6 +8,7 @@
 import numpy as np
 from numpy.core._rational_tests import rational
 from numpy.testing import assert_, assert_equal, assert_raises
+from numpy.core.numeric import pickle
 
 def assert_dtype_equal(a, b):
     assert_equal(a, b)

numpy/core/tests/test_multiarray.py

@@ -20,6 +20,9 @@
 import weakref
 import pytest
 from contextlib import contextmanager
+
+from numpy.core.numeric import pickle
+
 if sys.version_info[0] >= 3:
     import builtins
 else:
@@ -1371,7 +1374,6 @@ def test_view(self):
             assert_equal(zs.view((dt, 1)).shape, (0,))
 
     def test_pickle(self):
-        import pickle
         for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             for dt in [bytes, np.void, unicode]:
                 zs = self._zeros(10, dt)
@@ -3549,8 +3551,18 @@ def test_test_zero_rank(self):
 
 
 class TestPickling(object):
+    def test_highest_available_pickle_protocol(self):
+        try:
+            import pickle5
+        except ImportError:
+            pickle5 = None
+
+        if sys.version_info[:2] >= (3, 8) or pickle5 is not None:
+            assert pickle.HIGHEST_PROTOCOL >= 5
+        else:
+            assert pickle.HIGHEST_PROTOCOL < 5
+
     def test_roundtrip(self):
-        import pickle
         for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             carray = np.array([[2, 9], [7, 0], [3, 8]])
             DATA = [
@@ -3566,7 +3578,6 @@ def test_roundtrip(self):
                         err_msg="%r" % a)
 
     def _loads(self, obj):
-        import pickle
         if sys.version_info[0] >= 3:
             return pickle.loads(obj, encoding='latin1')
         else:

numpy/core/tests/test_overrides.py

@@ -1,6 +1,5 @@
 from __future__ import division, absolute_import, print_function
 
-import pickle
 import sys
 
 import numpy as np
@@ -9,6 +8,7 @@
 from numpy.core.overrides import (
     get_overloaded_types_and_args, array_function_dispatch,
     verify_matching_signatures)
+from numpy.core.numeric import pickle
 
 
 def _get_overloaded_args(relevant_args):
@@ -168,8 +168,10 @@ def dispatched_one_arg(array):
 class TestArrayFunctionDispatch(object):
 
     def test_pickle(self):
-        roundtripped = pickle.loads(pickle.dumps(dispatched_one_arg))
-        assert_(roundtripped is dispatched_one_arg)
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            roundtripped = pickle.loads(
+                    pickle.dumps(dispatched_one_arg, protocol=proto))
+            assert_(roundtripped is dispatched_one_arg)
 
     def test_name_and_docstring(self):
         assert_equal(dispatched_one_arg.__name__, 'dispatched_one_arg')

numpy/core/tests/test_records.py

@@ -7,7 +7,6 @@
     import collections.abc as collections_abc
 except ImportError:
     import collections as collections_abc
-import pickle
 import warnings
 import textwrap
 from os import path
@@ -18,6 +17,7 @@
     assert_, assert_equal, assert_array_equal, assert_array_almost_equal,
     assert_raises, assert_warns
     )
+from numpy.core.numeric import pickle
 
 
 class TestFromrecords(object):

numpy/core/tests/test_regression.py

@@ -1,7 +1,6 @@
 from __future__ import division, absolute_import, print_function
 
 import copy
-import pickle
 import sys
 import platform
 import gc
@@ -20,6 +19,7 @@
         _assert_valid_refcount, HAS_REFCOUNT,
         )
 from numpy.compat import asbytes, asunicode, long
+from numpy.core.numeric import pickle
 
 try:
     RecursionError
@@ -39,12 +39,13 @@ def test_mem_empty(self):
     def test_pickle_transposed(self):
         # Ticket #16
         a = np.transpose(np.array([[2, 9], [7, 0], [3, 8]]))
-        f = BytesIO()
-        pickle.dump(a, f)
-        f.seek(0)
-        b = pickle.load(f)
-        f.close()
-        assert_array_equal(a, b)
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            f = BytesIO()
+            pickle.dump(a, f, protocol=proto)
+            f.seek(0)
+            b = pickle.load(f)
+            f.close()
+            assert_array_equal(a, b)
 
     def test_typeNA(self):
         # Issue gh-515 
@@ -95,12 +96,13 @@ def test_negative_nd_indexing(self):
 
     def test_char_dump(self):
         # Ticket #50
-        f = BytesIO()
         ca = np.char.array(np.arange(1000, 1010), itemsize=4)
-        ca.dump(f)
-        f.seek(0)
-        ca = np.load(f)
-        f.close()
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            f = BytesIO()
+            pickle.dump(ca, f, protocol=proto)
+            f.seek(0)
+            ca = np.load(f)
+            f.close()
 
     def test_noncontiguous_fill(self):
         # Ticket #58.
@@ -359,12 +361,13 @@ def assign(a, b, c):
     def test_unpickle_dtype_with_object(self):
         # Implemented in r2840
         dt = np.dtype([('x', int), ('y', np.object_), ('z', 'O')])
-        f = BytesIO()
-        pickle.dump(dt, f)
-        f.seek(0)
-        dt_ = pickle.load(f)
-        f.close()
-        assert_equal(dt, dt_)
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            f = BytesIO()
+            pickle.dump(dt, f, protocol=proto)
+            f.seek(0)
+            dt_ = pickle.load(f)
+            f.close()
+            assert_equal(dt, dt_)
 
     def test_mem_array_creation_invalid_specification(self):
         # Ticket #196
@@ -474,7 +477,8 @@ def test_pickle_py2_bytes_encoding(self):
 
     def test_pickle_dtype(self):
         # Ticket #251
-        pickle.dumps(float)
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            pickle.dumps(float, protocol=proto)
 
     def test_swap_real(self):
         # Ticket #265

numpy/core/tests/test_ufunc.py

@@ -13,6 +13,7 @@
     assert_almost_equal, assert_array_almost_equal, assert_no_warnings,
     assert_allclose,
     )
+from numpy.core.numeric import pickle
 
 
 class TestUfuncKwargs(object):
@@ -43,16 +44,17 @@ def test_extobj_refcount(self):
 
 class TestUfunc(object):
     def test_pickle(self):
-        import pickle
-        assert_(pickle.loads(pickle.dumps(np.sin)) is np.sin)
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            assert_(pickle.loads(pickle.dumps(np.sin,
+                                              protocol=proto)) is np.sin)
 
-        # Check that ufunc not defined in the top level numpy namespace such as
-        # numpy.core._rational_tests.test_add can also be pickled
-        res = pickle.loads(pickle.dumps(_rational_tests.test_add))
-        assert_(res is _rational_tests.test_add)
+            # Check that ufunc not defined in the top level numpy namespace
+            # such as numpy.core._rational_tests.test_add can also be pickled
+            res = pickle.loads(pickle.dumps(_rational_tests.test_add,
+                                            protocol=proto))
+            assert_(res is _rational_tests.test_add)
 
     def test_pickle_withstring(self):
-        import pickle
         astring = (b"cnumpy.core\n_ufunc_reconstruct\np0\n"
                    b"(S'numpy.core.umath'\np1\nS'cos'\np2\ntp3\nRp4\n.")
         assert_(pickle.loads(astring) is np.cos)

numpy/lib/format.py

@@ -162,11 +162,8 @@
 import warnings
 from numpy.lib.utils import safe_eval
 from numpy.compat import asbytes, asstr, isfileobj, long, basestring
+from numpy.core.numeric import pickle
 
-if sys.version_info[0] >= 3:
-    import pickle
-else:
-    import cPickle as pickle
 
 MAGIC_PREFIX = b'\x93NUMPY'
 MAGIC_LEN = len(MAGIC_PREFIX) + 2

numpy/lib/npyio.py

@@ -23,12 +23,11 @@
     asbytes, asstr, asunicode, asbytes_nested, bytes, basestring, unicode,
     is_pathlib_path
     )
+from numpy.core.numeric import pickle
 
 if sys.version_info[0] >= 3:
-    import pickle
     from collections.abc import Mapping
 else:
-    import cPickle as pickle
     from future_builtins import map
     from collections import Mapping
 

numpy/ma/core.py

@@ -47,13 +47,9 @@
 from numpy.core.multiarray import normalize_axis_index
 from numpy.core.numeric import normalize_axis_tuple
 from numpy.core._internal import recursive
+from numpy.core.numeric import pickle
 
 
-if sys.version_info[0] >= 3:
-    import pickle
-else:
-    import cPickle as pickle
-
 __all__ = [
     'MAError', 'MaskError', 'MaskType', 'MaskedArray', 'abs', 'absolute',
     'add', 'all', 'allclose', 'allequal', 'alltrue', 'amax', 'amin',