Repository: numpy/numpy
File: numpy/_core/tests/test_multiarray.py (Python)

Code before:
        a = np.random.choice(10, 1000)
        b = np.random.choice(['abc', 'xy', 'wz', 'efghi', 'qwst', 'x'], 1000)

        for u in a, b:
            left = np.lexsort((u.astype('O'),))
            right = np.argsort(u, kind='mergesort')
            assert_array_equal(left, right)

        for u, v in (a, b), (b, a):
            idx = np.lexsort((u, v))
            assert_array_equal(idx, np.lexsort((u.astype('O'), v)))
            assert_array_equal(idx, np.lexsort((u, v.astype('O'))))
            u, v = np.array(u, dtype='object'), np.array(v, dtype='object')
            assert_array_equal(idx, np.lexsort((u, v)))

    def test_strings(self):  # gh-27984
        for dtype in "TU":
            surnames = np.array(['Hertz',    'Galilei', 'Hertz'], dtype=dtype)
            first_names = np.array(['Heinrich', 'Galileo', 'Gustav'], dtype=dtype)
            assert_array_equal(np.lexsort((first_names, surnames)), [1, 2, 0])

    def test_invalid_axis(self):  # gh-7528
        x = np.linspace(0., 1., 42 * 3).reshape(42, 3)
        assert_raises(AxisError, np.lexsort, x, axis=2)

@pytest.mark.thread_unsafe(reason="tmp_filename is thread-unsafe")
class TestIO:
    """Test tofile, fromfile, tobytes, and fromstring"""

    def _create_data(self):
        shape = (2, 4, 3)
        rand = np.random.random
        x = rand(shape) + rand(shape).astype(complex) * 1j
        x[0, :, 1] = [np.nan, np.inf, -np.inf, np.nan]
        return x

    @pytest.fixture(params=["string", "path_obj"])
    def tmp_filename(self, tmp_path, request):
        # This fixture covers two cases:
        # one where the filename is a string and
        # another where it is a pathlib object
        filename = tmp_path / "file"
        if request.param == "string":
            filename = str(filename)
        yield filename

    def test_nofile(self):
        # this should probably be supported as a file
        # but for now test for proper errors
        b = io.BytesIO()
        assert_raises(OSError, np.fromfile, b, np.uint8, 80)

Review comment (suggestion):
IMO we should just make the fixture thread-safe in numpy. You can do that by adding random data or unique thread-local data (e.g. the thread ID) to the temporary path.

Code after:
class TestLexsort:
    @pytest.mark.parametrize('dtype', [
        np.uint8, np.uint16, np.uint32, np.uint64,
        np.int8, np.int16, np.int32, np.int64,
        np.float16, np.float32, np.float64
    ])
    def test_basic(self, dtype):
        a = np.array([1, 2, 1, 3, 1, 5], dtype=dtype)
        b = np.array([0, 4, 5, 6, 2, 3], dtype=dtype)
        idx = np.lexsort((b, a))
        expected_idx = np.array([0, 4, 2, 1, 3, 5])
        assert_array_equal(idx, expected_idx)
        assert_array_equal(a[idx], np.sort(a))

    def test_mixed(self):
        a = np.array([1, 2, 1, 3, 1, 5])
        b = np.array([0, 4, 5, 6, 2, 3], dtype='datetime64[D]')

        idx = np.lexsort((b, a))
        expected_idx = np.array([0, 4, 2, 1, 3, 5])
        assert_array_equal(idx, expected_idx)

    def test_datetime(self):
        a = np.array([0, 0, 0], dtype='datetime64[D]')
        b = np.array([2, 1, 0], dtype='datetime64[D]')
        idx = np.lexsort((b, a))
        expected_idx = np.array([2, 1, 0])
        assert_array_equal(idx, expected_idx)

        a = np.array([0, 0, 0], dtype='timedelta64[D]')
        b = np.array([2, 1, 0], dtype='timedelta64[D]')
        idx = np.lexsort((b, a))
        expected_idx = np.array([2, 1, 0])
        assert_array_equal(idx, expected_idx)

    def test_object(self):  # gh-6312
        a = np.random.choice(10, 1000)
        b = np.random.choice(['abc', 'xy', 'wz', 'efghi', 'qwst', 'x'], 1000)

        for u in a, b:
            left = np.lexsort((u.astype('O'),))
            right = np.argsort(u, kind='mergesort')
            assert_array_equal(left, right)

        for u, v in (a, b), (b, a):
            idx = np.lexsort((u, v))
            assert_array_equal(idx, np.lexsort((u.astype('O'), v)))
            assert_array_equal(idx, np.lexsort((u, v.astype('O'))))
            u, v = np.array(u, dtype='object'), np.array(v, dtype='object')
            assert_array_equal(idx, np.lexsort((u, v)))