Fix lint and formatting errors - add missing npt import

Co-authored-by: talgalili <976006+talgalili@users.noreply.github.com>

Author: Copilot

balance/stats_and_plots/weighted_comparisons_stats.py

@@ -121,7 +121,9 @@ def _weights_per_covars_names(covar_names: List[str]) -> pd.DataFrame:
 
 
 def _kl_divergence_discrete(
-    p: npt.NDArray[np.floating[Any]], q: npt.NDArray[np.floating[Any]], eps: float = 1e-12
+    p: npt.NDArray[np.floating[Any]],
+    q: npt.NDArray[np.floating[Any]],
+    eps: float = 1e-12,
 ) -> float:
     """
     Compute the KL divergence between two discrete probability mass functions.

balance/util.py

@@ -30,6 +30,7 @@
 )
 
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 
 import pandas.api.types as pd_types

tests/test_adjustment.py

@@ -417,7 +417,9 @@ def test_apply_transformations_missing_column_specified(self) -> None:
     def test_apply_transformations_index_handling(self) -> None:
         """Test that dataframe indices are properly handled during transformations."""
         source_df = pd.DataFrame({"d": [1, 2, 3]}, index=(5, 6, 7))  # pyre-ignore[6]
-        target_df = pd.DataFrame({"d": [4, 5, 6, 7]}, index=(0, 1, 2, 3))  # pyre-ignore[6]
+        target_df = pd.DataFrame(
+            {"d": [4, 5, 6, 7]}, index=(0, 1, 2, 3)
+        )  # pyre-ignore[6]
         transformations = {"d": lambda x: x}
         result = apply_transformations((source_df, target_df), transformations)
         expected = (source_df, target_df)

tests/test_balancedf.py

@@ -253,7 +253,9 @@ def test_BalanceDFOutcomes_relative_response_rates(self) -> None:
         """
         self.assertEqual(
             s_o.outcomes().relative_response_rates(),
-            pd.DataFrame({"o1": [100.0, 4], "o2": [75.0, 3]}, index=["%", "n"]),  # pyre-ignore[6]
+            pd.DataFrame(
+                {"o1": [100.0, 4], "o2": [75.0, 3]}, index=["%", "n"]
+            ),  # pyre-ignore[6]
             lazy=True,
         )
 

tests/test_cbps.py

@@ -386,7 +386,9 @@ def test_cbps_consistency_with_default_arguments(self) -> None:
         # Create continuous variables for sample
         continuous_vars_sample = pd.concat(
             [
-                pd.DataFrame(np.random.uniform(0, 10, size=SAMPLE_SIZE), columns=[0]),  # pyre-ignore[6]
+                pd.DataFrame(
+                    np.random.uniform(0, 10, size=SAMPLE_SIZE), columns=[0]
+                ),  # pyre-ignore[6]
                 pd.DataFrame(
                     np.random.uniform(0, 1, size=(SAMPLE_SIZE, 4)), columns=range(1, 5)
                 ),

tests/test_rake.py

@@ -511,7 +511,9 @@ def test_rake_consistency_with_default_arguments(self) -> None:
         # Create sample DataFrame with mixed data types
         sample_df = pd.concat(
             [
-                pd.DataFrame(np.random.uniform(0, 10, size=n_sample), columns=[0]),  # pyre-ignore[6]
+                pd.DataFrame(
+                    np.random.uniform(0, 10, size=n_sample), columns=[0]
+                ),  # pyre-ignore[6]
                 pd.DataFrame(
                     np.random.uniform(0, 1, size=(n_sample, 4)), columns=range(1, 5)
                 ),
@@ -529,7 +531,9 @@ def test_rake_consistency_with_default_arguments(self) -> None:
         # Create target DataFrame with mixed data types
         target_df = pd.concat(
             [
-                pd.DataFrame(np.random.uniform(0, 18, size=n_target), columns=[0]),  # pyre-ignore[6]
+                pd.DataFrame(
+                    np.random.uniform(0, 18, size=n_target), columns=[0]
+                ),  # pyre-ignore[6]
                 pd.DataFrame(
                     np.random.uniform(0, 1, size=(n_target, 4)), columns=range(1, 5)
                 ),

tests/test_util.py

@@ -319,7 +319,9 @@ def test_drop_na_rows(self) -> None:
             sample_df,
             sample_weights,
         ) = balance_util.drop_na_rows(sample_df, sample_weights, "sample")
-        self.assertEqual(sample_df, pd.DataFrame({"a": (2.0), "b": ("c")}, index=[2]))  # pyre-ignore[6]
+        self.assertEqual(
+            sample_df, pd.DataFrame({"a": (2.0), "b": ("c")}, index=[2])
+        )  # pyre-ignore[6]
         self.assertEqual(sample_weights, pd.Series([3], index=[2]))
 
         # check exceptions
@@ -914,7 +916,9 @@ def test_quantize_preserves_column_order(self) -> None:
 
         result = balance_util.quantize(df, q=4, variables=["first", "third"])
 
-        self.assertListEqual(list(result.columns), ["first", "second", "third"])  # pyre-ignore[16]
+        self.assertListEqual(
+            list(result.columns), ["first", "second", "third"]
+        )  # pyre-ignore[16]
         self.assertIsInstance(result.loc[0, "first"], pd.Interval)  # pyre-ignore[16]
         self.assertEqual(result.loc[0, "second"], "a")  # pyre-ignore[16]
         self.assertIsInstance(result.loc[0, "third"], pd.Interval)  # pyre-ignore[16]
@@ -1582,7 +1586,9 @@ def test_fct_lump_by(self) -> None:
         )
 
         # test fct_lump_by doesn't affect indices when combining dataframes
-        s = pd.DataFrame({"d": [1, 1, 1], "e": ["a1", "a2", "a1"]}, index=(0, 6, 7))  # pyre-ignore[6]
+        s = pd.DataFrame(
+            {"d": [1, 1, 1], "e": ["a1", "a2", "a1"]}, index=(0, 6, 7)
+        )  # pyre-ignore[6]
         t = pd.DataFrame(
             {"d": [2, 3, 1, 2], "e": ["a2", "a2", "a1", "a2"]}, index=(0, 1, 2, 3)
         )