|
3 | 3 |
|
4 | 4 | from pandas import CategoricalIndex, DataFrame, Index, MultiIndex, Series, crosstab
|
5 | 5 | import pandas._testing as tm
|
| 6 | +from pandas.core.dtypes.common import is_categorical_dtype |
6 | 7 |
|
7 | 8 |
|
8 | 9 | class TestCrosstab:
|
@@ -743,3 +744,33 @@ def test_margin_normalize_multiple_columns(self):
|
743 | 744 | )
|
744 | 745 | expected.index.name = "C"
|
745 | 746 | tm.assert_frame_equal(result, expected)
|
| 747 | + |
| 748 | + |
| 749 | +@pytest.mark.parametrize("a_dtype", ["category", "int64"]) |
| 750 | +@pytest.mark.parametrize("b_dtype", ["category", "int64"]) |
| 751 | +def test_categoricals(a_dtype, b_dtype): |
| 752 | + # https://github.com/pandas-dev/pandas/issues/37465 |
| 753 | + g = np.random.RandomState(25982704) |
| 754 | + a = Series(g.randint(0, 3, size=100)).astype(a_dtype) |
| 755 | + b = Series(g.randint(0, 2, size=100)).astype(b_dtype) |
| 756 | + result = crosstab(a, b, margins=True, dropna=False) |
| 757 | + columns = Index([0, 1, "All"], dtype="object", name="col_0") |
| 758 | + index = Index([0, 1, 2, "All"], dtype="object", name="row_0") |
| 759 | + values = [[18, 16, 34], [18, 16, 34], [16, 16, 32], [52, 48, 100]] |
| 760 | + expected = DataFrame(values, index, columns) |
| 761 | + tm.assert_frame_equal(result, expected) |
| 762 | + |
| 763 | + # Verify when categorical does not have all values present |
| 764 | + a.loc[a == 1] = 2 |
| 765 | + a_is_cat = is_categorical_dtype(a.dtype) |
| 766 | + assert not a_is_cat or a.value_counts().loc[1] == 0 |
| 767 | + result = crosstab(a, b, margins=True, dropna=False) |
| 768 | + values = [[18, 16, 34], [0, 0, np.nan], [34, 32, 66], [52, 48, 100]] |
| 769 | + expected = DataFrame(values, index, columns) |
| 770 | + if not a_is_cat: |
| 771 | + expected = expected.loc[[0, 2, "All"]] |
| 772 | + expected["All"] = expected["All"].astype("int64") |
| 773 | + print(result) |
| 774 | + print(expected) |
| 775 | + print(expected.loc[[0, 2, "All"]]) |
| 776 | + tm.assert_frame_equal(result, expected) |
0 commit comments