format: make ty pass

Author: uriahf

src/rtichoke/_performance_table_spec.py

@@ -6,15 +6,15 @@
 
 from __future__ import annotations
 
+import math
 from collections.abc import Mapping
-from typing import Any, TypedDict
+from typing import Any, TypedDict, cast
 
-import math
 import polars as pl
 
 from rtichoke.processing.evaluation_semantics import _EvaluationMetadata
 
-_METRICS = (
+_METRICS: tuple[tuple[str, str], ...] = (
     ("true_positives", "True Positives"),
     ("true_negatives", "True Negatives"),
     ("false_positives", "False Positives"),
@@ -114,7 +114,7 @@ def _build_performance_table_spec(
         )
 
     metric_definitions: list[_MetricDefinition] = [
-        {"id": metric_id, "label": label}
+        cast(_MetricDefinition, {"id": metric_id, "label": label})
         for metric_id, label in _METRICS
         if metric_id in performance_data.columns
     ]
@@ -180,10 +180,16 @@ def _build_performance_table_spec(
         canonical_row: _PerformanceTableRow = {
             "evaluationId": evaluation_ids[group],
             "operatingPoint": operating_point,
-            "values": [
-                {"metricId": metric_id, "estimate": _nullable_number(row[metric_id])}
-                for metric_id in metric_ids
-            ],
+            "values": cast(
+                list[_MetricValue],
+                [
+                    {
+                        "metricId": metric_id,
+                        "estimate": _nullable_number(row[metric_id]),
+                    }
+                    for metric_id in metric_ids
+                ],
+            ),
         }
         if time_dependent:
             canonical_row["horizon"] = _number(row["fixed_time_horizon"])

src/rtichoke/processing/adjustments.py

@@ -1,8 +1,9 @@
+from collections.abc import Sequence
+
 import pandas as pd
 import polars as pl
-from polarstate import predict_aj_estimates
-from polarstate import prepare_event_table
-from collections.abc import Sequence
+from polarstate import predict_aj_estimates, prepare_event_table
+
 from rtichoke.processing.transforms import assign_and_explode_polars
 
 
@@ -529,7 +530,7 @@ def _aj_estimates_by_cutoff_per_horizon(
             df.filter(pl.col("fixed_time_horizon") == h)
             .group_by("strata")
             .map_groups(
-                lambda group: extract_aj_estimate_by_cutoffs(
+                lambda group, h=h: extract_aj_estimate_by_cutoffs(
                     group, [h], breaks, stratified_by, full_event_table=False
                 )
             )
@@ -547,7 +548,7 @@ def _aj_estimates_per_horizon(
             df.filter(pl.col("fixed_time_horizon") == h)
             .group_by("strata")
             .map_groups(
-                lambda group: extract_aj_estimate_for_strata(
+                lambda group, h=h: extract_aj_estimate_for_strata(
                     group, [h], full_event_table
                 )
             )
@@ -650,6 +651,8 @@ def _aj_adjusted_events(
             adjusted = extract_aj_estimate_by_cutoffs(
                 non_competing, horizons, breaks, stratified_by, full_event_table
             )
+        else:
+            raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")
 
         adjusted = adjusted.with_columns(
             [
@@ -704,6 +707,8 @@ def _aj_adjusted_events(
             adjusted = _aj_estimates_by_cutoff_per_horizon(
                 base_df, horizons, breaks, stratified_by
             )
+        else:
+            raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")
 
         adjusted = adjusted.with_columns(
             pl.lit(risk_set_scope)
@@ -730,6 +735,8 @@ def _aj_adjusted_events(
             adjusted = extract_aj_estimate_by_cutoffs(
                 base_df, horizons, breaks, stratified_by, full_event_table
             )
+        else:
+            raise ValueError(f"Unsupported risk-set scope: {risk_set_scope!r}")
 
         adjusted = adjusted.with_columns(
             [

tests/test_cutoff_grid_endpoint.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import numpy as np
 import pytest
 
@@ -17,7 +19,7 @@ def test_probability_threshold_breaks_match_r_seq_semantics(by, expected):
     breaks = create_breaks_values(None, "probability_threshold", by)
 
     assert breaks.tolist() == expected
-    assert np.max(breaks) <= 1.0
+    assert float(np.max(breaks)) <= 1.0
 
 
 @pytest.mark.parametrize(
@@ -33,7 +35,7 @@ def test_binary_performance_data_matches_r_cutoff_endpoint(by, expected_max):
 
     assert cutoffs.min() == 0.0
     assert cutoffs.max() == expected_max
-    assert cutoffs.max() <= 1.0
+    assert cast(float, cutoffs.max()) <= 1.0
 
 
 @pytest.mark.parametrize(
@@ -56,4 +58,4 @@ def test_time_performance_data_matches_r_cutoff_endpoint(by, expected_max):
 
     assert cutoffs.min() == 0.0
     assert cutoffs.max() == expected_max
-    assert cutoffs.max() <= 1.0
+    assert cast(float, cutoffs.max()) <= 1.0

tests/test_performance_table.py

@@ -1,18 +1,20 @@
+from typing import cast
+
 import numpy as np
 import pytest
 from great_tables import GT
 from reactable import Reactable
 
 import rtichoke.performance_table as performance_table_module
-from rtichoke.performance_table_reactable import _bar_style, _net_benefit_style
-
 from rtichoke import (
     create_performance_table,
     create_performance_table_times,
     prepare_performance_data,
     prepare_performance_data_times,
     render_performance_table,
 )
+from rtichoke.performance_table import PerformanceTableRenderer
+from rtichoke.performance_table_reactable import _bar_style, _net_benefit_style
 
 
 def _example():
@@ -215,7 +217,9 @@ def test_invalid_renderer_is_rejected():
     probs, reals = _example()
     data = prepare_performance_data(probs, reals, by=0.1)
     with pytest.raises(ValueError, match="renderer"):
-        render_performance_table(data, renderer="unknown")
+        render_performance_table(
+            data, renderer=cast(PerformanceTableRenderer, "unknown")
+        )
 
 
 def test_reactable_metric_bar_matches_r_colors_and_geometry():

tests/test_performance_table_spec.py

@@ -21,9 +21,7 @@ def _static(
     by: float = 0.5,
     stratified_by: Sequence[str] = ("probability_threshold",),
 ):
-    data = prepare_performance_data(
-        probs, reals, by=by, stratified_by=stratified_by
-    )
+    data = prepare_performance_data(probs, reals, by=by, stratified_by=stratified_by)
     metadata = _build_evaluation_metadata(probs, reals, np.array([]))
     return _performance_table_spec_from_performance_data(data, metadata)
 
@@ -110,9 +108,7 @@ def test_static_probability_threshold_and_ppcr_operating_points():
         "probability_threshold"
     }
     assert {row["operatingPoint"]["type"] for row in ppcr["rows"]} == {"ppcr"}
-    assert all(
-        0 <= row["operatingPoint"]["value"] <= 1 for row in ppcr["rows"]
-    )
+    assert all(0 <= row["operatingPoint"]["value"] <= 1 for row in ppcr["rows"])
 
 
 def test_zero_is_preserved_and_missing_metric_is_null():

tests/test_time_gains_v2.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Any, cast
 
 import matplotlib.figure
 import numpy as np
@@ -116,6 +117,7 @@ def test_censoring_and_competing_risk_reference_comes_from_performance_layer():
     spec = _gains_times_v2_spec_from_performance_data(
         performance, _build_evaluation_metadata(probs, reals, times)
     )
+    references = cast(list[dict[str, Any]], spec["references"])
     calculated = {
         float(row["fixed_time_horizon"]): float(row["real_positives"] / row["n"])
         for row in performance.filter(performance["chosen_cutoff"] == 0)
@@ -126,7 +128,7 @@ def test_censoring_and_competing_risk_reference_comes_from_performance_layer():
 
     assert {
         reference["horizon"]: reference["points"][1]["x"]
-        for reference in spec["references"][1:]
+        for reference in references[1:]
     } == calculated
 
 

tests/test_time_lift_v2.py

@@ -1,3 +1,5 @@
+from typing import Any, cast
+
 import matplotlib.figure
 import numpy as np
 import plotly.graph_objects as go
@@ -143,6 +145,7 @@ def test_time_lift_censoring_and_competing_risk_reference_comes_from_performance
     spec = _lift_times_v2_spec_from_performance_data(
         performance, _build_evaluation_metadata(probs, reals, times)
     )
+    references = cast(list[dict[str, Any]], spec["references"])
     calculated_risks = {
         float(row["fixed_time_horizon"]): float(row["real_positives"] / row["n"])
         for row in performance.filter(performance["chosen_cutoff"] == 0)
@@ -151,7 +154,7 @@ def test_time_lift_censoring_and_competing_risk_reference_comes_from_performance
         .to_dicts()
     }
 
-    for reference in spec["references"][1:]:
+    for reference in references[1:]:
         horizon = reference["horizon"]
         risk = calculated_risks[horizon]
         assert reference["points"] == [

tests/test_viz_browser.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Any, cast
 
 import numpy as np
 
@@ -35,6 +36,7 @@ def test_real_roc_output_maps_to_canonical_spec():
     performance_data = _real_roc_performance_data()
 
     spec = _roc_spec_from_performance_data(performance_data)
+    spec = cast(dict[str, Any], spec)
 
     assert spec["schemaVersion"] == "1.0"
     assert spec["type"] == "roc"
@@ -47,6 +49,7 @@ def test_real_roc_output_maps_to_canonical_spec():
 
 def test_real_calibration_output_maps_to_canonical_spec():
     spec = _calibration_spec_from_curve_list(_real_calibration_curve_list())
+    spec = cast(dict[str, Any], spec)
 
     assert spec["schemaVersion"] == "1.0"
     assert spec["type"] == "calibration"

tests/test_viz_spec_v2.py

@@ -1,3 +1,5 @@
+from typing import Any, cast
+
 import numpy as np
 
 from rtichoke._viz_spec_v2 import (
@@ -24,6 +26,7 @@ def test_roc_v2_one_model_one_population():
         performance_data,
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     assert spec["schemaVersion"] == "2.0"
     assert spec["type"] == "roc"
@@ -57,6 +60,7 @@ def test_roc_v2_multiple_models_share_one_population():
         performance_data,
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     evaluations = spec["evaluations"]
     assert [evaluation["id"] for evaluation in evaluations] == [
@@ -87,6 +91,7 @@ def test_roc_v2_keyed_populations_keep_model_identity_unknown():
         performance_data,
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     assert spec["evaluations"] == [
         {"id": "evaluation-1", "population": "Population A"},
@@ -112,6 +117,7 @@ def test_roc_v2_ids_do_not_encode_compatibility_group_labels():
         prepare_performance_data(probs, reals, by=0.25),
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     renamed_reals = {
         "Cohort X": reals["Population A"],
@@ -125,6 +131,7 @@ def test_roc_v2_ids_do_not_encode_compatibility_group_labels():
         prepare_performance_data(renamed_probs, renamed_reals, by=0.25),
         _static_metadata(renamed_probs, renamed_reals),
     )
+    renamed_spec = cast(dict[str, Any], renamed_spec)
 
     assert [evaluation["id"] for evaluation in spec["evaluations"]] == [
         evaluation["id"] for evaluation in renamed_spec["evaluations"]
@@ -142,6 +149,7 @@ def test_gains_v2_uses_production_prevalence_for_perfect_path():
     spec = _gains_v2_spec_from_performance_data(
         performance_data, _static_metadata(probs, reals)
     )
+    spec = cast(dict[str, Any], spec)
 
     assert spec["type"] == "gains"
     assert spec["x"] == "ppcr"
@@ -171,6 +179,7 @@ def test_gains_v2_shares_one_perfect_path_across_models():
         prepare_performance_data(probs, reals, by=0.25),
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     assert len(spec["series"]) == 2
     assert len(spec["references"]) == 2
@@ -191,6 +200,7 @@ def test_gains_v2_keeps_equal_prevalence_populations_distinct():
         prepare_performance_data(probs, reals, by=0.25),
         _static_metadata(probs, reals),
     )
+    spec = cast(dict[str, Any], spec)
 
     perfect = spec["references"][1:]
     assert [reference["population"] for reference in perfect] == [