1import numpy as np
2import pytest
3
4from pandas import Index, date_range
5import pandas._testing as tm
6from pandas.core.reshape.util import cartesian_product
7
8
9class TestCartesianProduct:
10    def test_simple(self):
11        x, y = list("ABC"), [1, 22]
12        result1, result2 = cartesian_product([x, y])
13        expected1 = np.array(["A", "A", "B", "B", "C", "C"])
14        expected2 = np.array([1, 22, 1, 22, 1, 22])
15        tm.assert_numpy_array_equal(result1, expected1)
16        tm.assert_numpy_array_equal(result2, expected2)
17
18    def test_datetimeindex(self):
19        # regression test for GitHub issue #6439
20        # make sure that the ordering on datetimeindex is consistent
21        x = date_range("2000-01-01", periods=2)
22        result1, result2 = [Index(y).day for y in cartesian_product([x, x])]
23        expected1 = Index([1, 1, 2, 2])
24        expected2 = Index([1, 2, 1, 2])
25        tm.assert_index_equal(result1, expected1)
26        tm.assert_index_equal(result2, expected2)
27
28    def test_tzaware_retained(self):
29        x = date_range("2000-01-01", periods=2, tz="US/Pacific")
30        y = np.array([3, 4])
31        result1, result2 = cartesian_product([x, y])
32
33        expected = x.repeat(2)
34        tm.assert_index_equal(result1, expected)
35
36    def test_tzaware_retained_categorical(self):
37        x = date_range("2000-01-01", periods=2, tz="US/Pacific").astype("category")
38        y = np.array([3, 4])
39        result1, result2 = cartesian_product([x, y])
40
41        expected = x.repeat(2)
42        tm.assert_index_equal(result1, expected)
43
44    def test_empty(self):
45        # product of empty factors
46        X = [[], [0, 1], []]
47        Y = [[], [], ["a", "b", "c"]]
48        for x, y in zip(X, Y):
49            expected1 = np.array([], dtype=np.asarray(x).dtype)
50            expected2 = np.array([], dtype=np.asarray(y).dtype)
51            result1, result2 = cartesian_product([x, y])
52            tm.assert_numpy_array_equal(result1, expected1)
53            tm.assert_numpy_array_equal(result2, expected2)
54
55        # empty product (empty input):
56        result = cartesian_product([])
57        expected = []
58        assert result == expected
59
60    @pytest.mark.parametrize(
61        "X", [1, [1], [1, 2], [[1], 2], "a", ["a"], ["a", "b"], [["a"], "b"]]
62    )
63    def test_invalid_input(self, X):
64        msg = "Input must be a list-like of list-likes"
65
66        with pytest.raises(TypeError, match=msg):
67            cartesian_product(X=X)
68
69    def test_exceed_product_space(self):
70        # GH31355: raise useful error when produce space is too large
71        msg = "Product space too large to allocate arrays!"
72
73        with pytest.raises(ValueError, match=msg):
74            dims = [np.arange(0, 22, dtype=np.int16) for i in range(12)] + [
75                (np.arange(15128, dtype=np.int16)),
76            ]
77            cartesian_product(X=dims)
78