Preserve bool dtype in SingleAmountTaxScale.calc()

changelog.d/bool-bracket-dtype.fixed.md

@@ -0,0 +1 @@
+`SingleAmountTaxScale.calc()` now preserves the bracket amounts' dtype, so boolean `single_amount` brackets (`amount_unit: bool`) return a boolean array instead of an int64 `0`/`1` array; this restores logical NOT (`~`) on the result, which previously performed a bitwise negation and silently produced wrong values.

policyengine_core/taxscales/single_amount_tax_scale.py

@@ -19,6 +19,15 @@ def calc(
         """
         Matches the input amount to a set of brackets and returns the single
         cell value that fits within that bracket.
+
+        The returned array preserves the dtype of the bracket amounts. In
+        particular, a scale whose amounts are booleans (``amount_unit: bool``
+        with ``amount: true``/``false``) returns a boolean array rather than the
+        int64 ``0``/``1`` array that naive sentinel padding would produce. This
+        keeps logical NOT (``~``) working on the result: ``~`` on an int array
+        is a bitwise negation (``~1 == -2`` and ``~0 == -1``, both truthy),
+        whereas ``~`` on a bool array is the intended element-wise logical
+        negation.
         """
         guarded_thresholds = numpy.array([-numpy.inf] + self.thresholds + [numpy.inf])
 
@@ -28,6 +37,18 @@ def calc(
             right=right,
         )
 
-        guarded_amounts = numpy.array([0] + self.amounts + [0])
+        # Build the amounts array from the bracket amounts first, so numpy
+        # infers their natural dtype (bool for boolean brackets, int/float
+        # otherwise), then pad the two out-of-range sentinel cells with a zero
+        # of that same dtype (``False`` for bool, ``0`` for numeric). Padding
+        # with a Python int ``0`` up front, as the previous implementation did,
+        # coerced the whole array to int64 and silently dropped a boolean dtype.
+        # With no amounts, fall back to the historical int64 zero array.
+        if len(self.amounts) == 0:
+            guarded_amounts = numpy.array([0, 0])
+        else:
+            amounts = numpy.asarray(self.amounts)
+            sentinel = numpy.zeros(1, dtype=amounts.dtype)
+            guarded_amounts = numpy.concatenate([sentinel, amounts, sentinel])
 
         return guarded_amounts[bracket_indices - 1]

tests/core/tax_scales/test_single_amount_tax_scale.py

@@ -45,6 +45,117 @@ def test_to_dict():
     assert result == {"6": 0.23, "9": 0.29}
 
 
+def test_calc_preserves_boolean_dtype():
+    # A single_amount scale whose amounts are booleans (amount_unit: bool)
+    # must return a boolean array, not an int64 0/1 array. Otherwise logical
+    # NOT (~) silently becomes a bitwise negation on the result.
+    tax_scale = taxscales.SingleAmountTaxScale()
+    tax_scale.add_bracket(0, True)
+    tax_scale.add_bracket(16, False)
+
+    result = tax_scale.calc(numpy.array([10.0, 20.0]))
+
+    assert result.dtype == numpy.bool_
+    assert result.tolist() == [True, False]
+
+
+def test_calc_boolean_result_supports_logical_not():
+    # The regression this guards: ~ on the int64 output was a bitwise negation
+    # (~1 == -2, ~0 == -1, both truthy), so negating a bool bracket produced
+    # wrong results. On a real bool array, ~ is the intended logical NOT.
+    tax_scale = taxscales.SingleAmountTaxScale()
+    tax_scale.add_bracket(0, True)
+    tax_scale.add_bracket(16, False)
+
+    result = tax_scale.calc(numpy.array([10.0, 20.0]))
+    negated = ~result
+
+    assert negated.dtype == numpy.bool_
+    assert negated.tolist() == [False, True]
+
+
+def test_calc_preserves_integer_dtype():
+    # Numeric int amounts must be unaffected by the boolean-dtype fix: the
+    # result stays int64 with identical values.
+    tax_scale = taxscales.SingleAmountTaxScale()
+    tax_scale.add_bracket(0, 100)
+    tax_scale.add_bracket(16, 200)
+
+    result = tax_scale.calc(numpy.array([10.0, 20.0]))
+
+    assert result.dtype == numpy.int64
+    assert result.tolist() == [100, 200]
+
+
+def test_calc_preserves_float_dtype():
+    # Numeric float amounts must be unaffected: the result stays float64.
+    tax_scale = taxscales.SingleAmountTaxScale()
+    tax_scale.add_bracket(6, 0.23)
+    tax_scale.add_bracket(9, 0.29)
+
+    result = tax_scale.calc(numpy.array([1, 8, 10]))
+
+    assert result.dtype == numpy.float64
+    tools.assert_near(result, [0, 0.23, 0.29])
+
+
+def test_calc_out_of_range_bool_returns_false_sentinel():
+    # Inputs that fall outside every bracket must get the zero-valued sentinel
+    # in the amounts' own dtype: False for a boolean scale, not int 0.
+    tax_scale = taxscales.SingleAmountTaxScale()
+    # digitize guards with [-inf, threshold..., inf]; the -inf..first-threshold
+    # cell is a real bracket, so use a scale where only interior cells map to a
+    # bracket and confirm the sentinel dtype anyway via a below-all base.
+    tax_scale.add_bracket(0, True)
+
+    result = tax_scale.calc(numpy.array([-5.0]))
+
+    assert result.dtype == numpy.bool_
+
+
+def test_calc_empty_scale_stays_int():
+    # An empty scale has no amounts to infer a dtype from; preserve the
+    # historical behaviour of returning an int64 zero array.
+    tax_scale = taxscales.SingleAmountTaxScale()
+
+    result = tax_scale.calc(numpy.array([10.0]))
+
+    assert result.dtype == numpy.int64
+    assert result.tolist() == [0]
+
+
+def test_calc_boolean_dtype_end_to_end_via_parameter_scale():
+    # Mirror how policyengine-us declares boolean single_amount brackets
+    # (amount_unit: bool, amount: true/false) and confirm the dtype survives
+    # the full ParameterScale -> get_at_instant -> SingleAmountTaxScale path.
+    data = {
+        "description": "Boolean age-exemption single_amount scale",
+        "metadata": {
+            "type": "single_amount",
+            "threshold_unit": "year",
+            "amount_unit": "bool",
+        },
+        "brackets": [
+            {
+                "threshold": {"2022-01-01": {"value": 0}},
+                "amount": {"2022-01-01": {"value": True}},
+            },
+            {
+                "threshold": {"2022-01-01": {"value": 16}},
+                "amount": {"2022-01-01": {"value": False}},
+            },
+        ],
+    }
+    scale = parameters.ParameterScale("bool_scale", data, "")
+    scale_at_instant = scale.get_at_instant(periods.Instant((2022, 6, 1)))
+
+    result = scale_at_instant.calc(numpy.array([10.0, 20.0]))
+
+    assert result.dtype == numpy.bool_
+    assert result.tolist() == [True, False]
+    assert (~result).tolist() == [False, True]
+
+
 # TODO: move, as we're testing Scale, not SingleAmountTaxScale
 def test_assign_thresholds_on_creation(data):
     scale = parameters.ParameterScale("amount_scale", data, "")