diff --git a/patsy/categorical.py b/patsy/categorical.py index c88ecd5..0eb772e 100644 --- a/patsy/categorical.py +++ b/patsy/categorical.py @@ -360,7 +360,7 @@ def categorical_to_int(data, levels, NA_action, origin=None): # fastpath to avoid doing an item-by-item iteration over boolean arrays, # as requested by #44 if hasattr(data, "dtype") and safe_issubdtype(data.dtype, np.bool_): - if level_to_int[False] == 0 and level_to_int[True] == 1: + if level_to_int.get(False) == 0 and level_to_int.get(True) == 1: return data.astype(np.int_) out = np.empty(len(data), dtype=int) for i, value in enumerate(data): diff --git a/patsy/test_regressions.py b/patsy/test_regressions.py index 2760846..fc9b807 100644 --- a/patsy/test_regressions.py +++ b/patsy/test_regressions.py @@ -6,6 +6,9 @@ # else) from patsy import EvalEnvironment, dmatrix, build_design_matrices, PatsyError, Origin +import numpy as np +import pytest +from patsy.util import have_pandas def test_issue_11(): @@ -22,3 +25,45 @@ def test_issue_11(): assert e.origin == Origin(formula, 0, 4) else: assert False + + +@pytest.mark.parametrize("level", [False, True]) +def test_single_boolean_level_array(level): + matrix = dmatrix( + "0 + C(x, levels=levels)", + {"x": np.array([level, level]), "levels": [level]}, + ) + np.testing.assert_array_equal(matrix, [[1], [1]]) + info = matrix.design_info + assert next(iter(info.factor_infos.values())).categories == (level,) + prediction = build_design_matrices( + [info], {"x": np.array([level]), "levels": [level]} + )[0] + np.testing.assert_array_equal(prediction, [[1]]) + + +@pytest.mark.parametrize("level", [False, True]) +def test_single_boolean_level_prediction(level): + matrix = dmatrix("0 + C(x, levels=levels)", {"x": [level], "levels": [level]}) + prediction = build_design_matrices( + [matrix.design_info], {"x": np.array([level]), "levels": [level]} + )[0] + np.testing.assert_array_equal(prediction, [[1]]) + with pytest.raises(PatsyError, match="does not match any of the expected levels"): + build_design_matrices( + [matrix.design_info], {"x": np.array([not level]), "levels": [level]} + ) + + +@pytest.mark.skipif(not have_pandas, reason="pandas is not available") +@pytest.mark.parametrize("level", [False, True]) +def test_single_boolean_level_series(level): + import pandas + + matrix = dmatrix( + "0 + C(x, levels=levels)", + {"x": pandas.Series([level, level], index=[10, 20]), "levels": [level]}, + return_type="dataframe", + ) + np.testing.assert_array_equal(matrix, [[1], [1]]) + assert list(matrix.index) == [10, 20]