12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879 |
- import numpy as np
- import pytest
- from pandas import (
- Index,
- date_range,
- )
- import pandas._testing as tm
- from pandas.core.reshape.util import cartesian_product
- class TestCartesianProduct:
- def test_simple(self):
- x, y = list("ABC"), [1, 22]
- result1, result2 = cartesian_product([x, y])
- expected1 = np.array(["A", "A", "B", "B", "C", "C"])
- expected2 = np.array([1, 22, 1, 22, 1, 22])
- tm.assert_numpy_array_equal(result1, expected1)
- tm.assert_numpy_array_equal(result2, expected2)
- def test_datetimeindex(self):
- # regression test for GitHub issue #6439
- # make sure that the ordering on datetimeindex is consistent
- x = date_range("2000-01-01", periods=2)
- result1, result2 = (Index(y).day for y in cartesian_product([x, x]))
- expected1 = Index([1, 1, 2, 2], dtype=np.int32)
- expected2 = Index([1, 2, 1, 2], dtype=np.int32)
- tm.assert_index_equal(result1, expected1)
- tm.assert_index_equal(result2, expected2)
- def test_tzaware_retained(self):
- x = date_range("2000-01-01", periods=2, tz="US/Pacific")
- y = np.array([3, 4])
- result1, result2 = cartesian_product([x, y])
- expected = x.repeat(2)
- tm.assert_index_equal(result1, expected)
- def test_tzaware_retained_categorical(self):
- x = date_range("2000-01-01", periods=2, tz="US/Pacific").astype("category")
- y = np.array([3, 4])
- result1, result2 = cartesian_product([x, y])
- expected = x.repeat(2)
- tm.assert_index_equal(result1, expected)
- @pytest.mark.parametrize("x, y", [[[], []], [[0, 1], []], [[], ["a", "b", "c"]]])
- def test_empty(self, x, y):
- # product of empty factors
- expected1 = np.array([], dtype=np.asarray(x).dtype)
- expected2 = np.array([], dtype=np.asarray(y).dtype)
- result1, result2 = cartesian_product([x, y])
- tm.assert_numpy_array_equal(result1, expected1)
- tm.assert_numpy_array_equal(result2, expected2)
- def test_empty_input(self):
- # empty product (empty input):
- result = cartesian_product([])
- expected = []
- assert result == expected
- @pytest.mark.parametrize(
- "X", [1, [1], [1, 2], [[1], 2], "a", ["a"], ["a", "b"], [["a"], "b"]]
- )
- def test_invalid_input(self, X):
- msg = "Input must be a list-like of list-likes"
- with pytest.raises(TypeError, match=msg):
- cartesian_product(X=X)
- def test_exceed_product_space(self):
- # GH31355: raise useful error when produce space is too large
- msg = "Product space too large to allocate arrays!"
- with pytest.raises(ValueError, match=msg):
- dims = [np.arange(0, 22, dtype=np.int16) for i in range(12)] + [
- (np.arange(15128, dtype=np.int16)),
- ]
- cartesian_product(X=dims)
|