From 69e53f45b57f7b1bc19fc897195f3a9ca826023d Mon Sep 17 00:00:00 2001
From: Valentin Laurent <valentin.laurent.fr@gmail.com>
Date: Tue, 4 Feb 2025 12:06:07 +0100
Subject: [PATCH] CHORE: make type checking stricter for v1 files (#609)

---
 Makefile                                       |  2 +-
 mapie_v1/_utils.py                             |  2 +-
 .../integration_tests/tests/test_regression.py | 18 +++++++++++-------
 mapie_v1/integration_tests/utils.py            |  9 +++++----
 4 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/Makefile b/Makefile
index 24e4656a8..f471b4614 100644
--- a/Makefile
+++ b/Makefile
@@ -9,7 +9,7 @@ type-check:
 	mypy mapie
 
 v1-type-check:
-	mypy mapie_v1 --exclude $(mapie_v0_folder_name)
+	mypy mapie_v1 --disallow-untyped-defs --exclude $(mapie_v0_folder_name)
 
 tests:
 	pytest -vs --doctest-modules mapie
diff --git a/mapie_v1/_utils.py b/mapie_v1/_utils.py
index 06caa7b42..4a0496087 100644
--- a/mapie_v1/_utils.py
+++ b/mapie_v1/_utils.py
@@ -34,7 +34,7 @@ def check_if_param_in_allowed_values(
         )
 
 
-def check_cv_not_string(cv: Union[int, str, BaseCrossValidator]):
+def check_cv_not_string(cv: Union[int, str, BaseCrossValidator]) -> None:
     if isinstance(cv, str):
         raise ValueError(
             "'cv' string options not available in MAPIE >= v1"
diff --git a/mapie_v1/integration_tests/tests/test_regression.py b/mapie_v1/integration_tests/tests/test_regression.py
index 4d48c7495..8dea54ad1 100644
--- a/mapie_v1/integration_tests/tests/test_regression.py
+++ b/mapie_v1/integration_tests/tests/test_regression.py
@@ -12,7 +12,7 @@
 from sklearn.model_selection import train_test_split
 
 from mapie.subsample import Subsample
-from mapie._typing import ArrayLike
+from mapie._typing import ArrayLike, NDArray
 from mapie.conformity_scores import GammaConformityScore, \
     AbsoluteConformityScore, ResidualNormalisedScore
 from mapie_v1.regression import SplitConformalRegressor, \
@@ -136,7 +136,7 @@
 
 
 @pytest.mark.parametrize("params_split", params_test_cases_split)
-def test_intervals_and_predictions_exact_equality_split(params_split):
+def test_intervals_and_predictions_exact_equality_split(params_split: dict) -> None:
     v0_params = params_split["v0"]
     v1_params = params_split["v1"]
 
@@ -222,7 +222,7 @@ def test_intervals_and_predictions_exact_equality_split(params_split):
 
 
 @pytest.mark.parametrize("params_cross", params_test_cases_cross)
-def test_intervals_and_predictions_exact_equality_cross(params_cross):
+def test_intervals_and_predictions_exact_equality_cross(params_cross: dict) -> None:
 
     compare_model_predictions_and_intervals(
         model_v0=MapieRegressorV0,
@@ -313,7 +313,9 @@ def test_intervals_and_predictions_exact_equality_cross(params_cross):
 
 
 @pytest.mark.parametrize("params_jackknife", params_test_cases_jackknife)
-def test_intervals_and_predictions_exact_equality_jackknife(params_jackknife):
+def test_intervals_and_predictions_exact_equality_jackknife(
+    params_jackknife: dict
+) -> None:
 
     compare_model_predictions_and_intervals(
         model_v0=MapieRegressorV0,
@@ -419,7 +421,9 @@ def test_intervals_and_predictions_exact_equality_jackknife(params_jackknife):
 
 
 @pytest.mark.parametrize("params_quantile", params_test_cases_quantile)
-def test_intervals_and_predictions_exact_equality_quantile(params_quantile):
+def test_intervals_and_predictions_exact_equality_quantile(
+    params_quantile: dict
+) -> None:
     v0_params = params_quantile["v0"]
     v1_params = params_quantile["v1"]
 
@@ -447,8 +451,8 @@ def compare_model_predictions_and_intervals(
         JackknifeAfterBootstrapRegressor,
         ConformalizedQuantileRegressor
     ]],
-    X: ArrayLike,
-    y: ArrayLike,
+    X: NDArray,
+    y: NDArray,
     v0_params: Dict = {},
     v1_params: Dict = {},
     prefit: bool = False,
diff --git a/mapie_v1/integration_tests/utils.py b/mapie_v1/integration_tests/utils.py
index 0634c9efe..bb816efed 100644
--- a/mapie_v1/integration_tests/utils.py
+++ b/mapie_v1/integration_tests/utils.py
@@ -1,14 +1,15 @@
-from typing import Callable, Dict, Any, Optional
+from typing import Callable, Dict, Any, Optional, Tuple
+from mapie._typing import NDArray
 import inspect
 from sklearn.model_selection import ShuffleSplit
 
 
 def train_test_split_shuffle(
-    X,
-    y,
+    X: NDArray,
+    y: NDArray,
     test_size: float = 0.2,
     random_state: int = 42
-):
+) -> Tuple[Any, Any, Any, Any]:
 
     splitter = ShuffleSplit(n_splits=1,
                             test_size=test_size,