Add Python type hint for tests and demos. (#11795)

Author: trivialfis

demo/guide-python/cross_validation.py

@@ -4,6 +4,7 @@
 """
 
 import os
+from typing import Any, Dict, Tuple
 
 import numpy as np
 
@@ -54,7 +55,9 @@
 # used to return the preprocessed training, test data, and parameter
 # we can use this to do weight rescale, etc.
 # as a example, we try to set scale_pos_weight
-def fpreproc(dtrain, dtest, param):
+def fpreproc(
+    dtrain: xgb.DMatrix, dtest: xgb.DMatrix, param: Any
+) -> Tuple[xgb.DMatrix, xgb.DMatrix, Dict[str, Any]]:
     label = dtrain.get_label()
     ratio = float(np.sum(label == 0)) / np.sum(label == 1)
     param["scale_pos_weight"] = ratio
@@ -74,15 +77,15 @@ def fpreproc(dtrain, dtest, param):
 print("running cross validation, with customized loss function")
 
 
-def logregobj(preds, dtrain):
+def logregobj(preds: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[np.ndarray, np.ndarray]:
     labels = dtrain.get_label()
     preds = 1.0 / (1.0 + np.exp(-preds))
     grad = preds - labels
     hess = preds * (1.0 - preds)
     return grad, hess
 
 
-def evalerror(preds, dtrain):
+def evalerror(preds: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, float]:
     labels = dtrain.get_label()
     preds = 1.0 / (1.0 + np.exp(-preds))
     return "error", float(sum(labels != (preds > 0.0))) / len(labels)

demo/guide-python/custom_rmsle.py

@@ -18,7 +18,6 @@
 from time import time
 from typing import Dict, List, Tuple
 
-import matplotlib
 import numpy as np
 from matplotlib import pyplot as plt
 
@@ -136,7 +135,7 @@ def squared_log(predt: np.ndarray,
     def rmsle(predt: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, float]:
         ''' Root mean squared log error metric.
 
-        :math:`\sqrt{\frac{1}{N}[log(pred + 1) - log(label + 1)]^2}`
+        :math:`\\sqrt{\frac{1}{N}[log(pred + 1) - log(label + 1)]^2}`
         '''
         y = dtrain.get_label()
         predt[predt < -1] = -1 + 1e-6
@@ -156,11 +155,16 @@ def rmsle(predt: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, float]:
     return results
 
 
-def plot_history(rmse_evals, rmsle_evals, py_rmsle_evals):
+def plot_history(
+    rmse_evals: Dict[str, Dict],
+    rmsle_evals: Dict[str, Dict],
+    py_rmsle_evals: Dict[str, Dict]
+) -> None:
     fig, axs = plt.subplots(3, 1)
-    ax0: matplotlib.axes.Axes = axs[0]
-    ax1: matplotlib.axes.Axes = axs[1]
-    ax2: matplotlib.axes.Axes = axs[2]
+    assert isinstance(axs, np.ndarray)
+    ax0 = axs[0]
+    ax1 = axs[1]
+    ax2 = axs[2]
 
     x = np.arange(0, kBoostRound, 1)
 
@@ -177,7 +181,7 @@ def plot_history(rmse_evals, rmsle_evals, py_rmsle_evals):
     ax2.legend()
 
 
-def main(args):
+def main(args: argparse.Namespace) -> None:
     dtrain, dtest = generate_data()
     rmse_evals = native_rmse(dtrain, dtest)
     rmsle_evals = native_rmsle(dtrain, dtest)

demo/guide-python/evals_result.py

@@ -3,6 +3,7 @@
 ======================================================
 """
 import os
+from typing import Any, Dict
 
 import xgboost as xgb
 
@@ -24,7 +25,7 @@
 num_round = 2
 watchlist = [(dtest, "eval"), (dtrain, "train")]
 
-evals_result = {}
+evals_result: Dict[str, Any] = {}
 bst = xgb.train(param, dtrain, num_round, watchlist, evals_result=evals_result)
 
 print("Access logloss metric directly from evals_result:")

demo/guide-python/predict_first_ntree.py

@@ -14,7 +14,7 @@
 test = os.path.join(CURRENT_DIR, "../data/agaricus.txt.test")
 
 
-def native_interface():
+def native_interface() -> None:
     # load data in do training
     dtrain = xgb.DMatrix(train + "?format=libsvm")
     dtest = xgb.DMatrix(test + "?format=libsvm")
@@ -34,7 +34,7 @@ def native_interface():
     print("error of ypred2=%f" % (np.sum((ypred2 > 0.5) != label) / float(len(label))))
 
 
-def sklearn_interface():
+def sklearn_interface() -> None:
     X_train, y_train = load_svmlight_file(train)
     X_test, y_test = load_svmlight_file(test)
     clf = xgb.XGBClassifier(n_estimators=3, max_depth=2, eta=1)

demo/guide-python/spark_estimator_examples.py

@@ -5,10 +5,11 @@
 @author: Weichen Xu
 """
 
+import numpy as np
 import sklearn.datasets
 from pyspark.ml.evaluation import MulticlassClassificationEvaluator, RegressionEvaluator
 from pyspark.ml.linalg import Vectors
-from pyspark.sql import SparkSession
+from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.functions import rand
 from sklearn.model_selection import train_test_split
 
@@ -17,7 +18,7 @@
 spark = SparkSession.builder.master("local[*]").getOrCreate()
 
 
-def create_spark_df(X, y):
+def create_spark_df(X: np.ndarray, y: np.ndarray) -> DataFrame:
     return spark.createDataFrame(
         spark.sparkContext.parallelize(
             [(Vectors.dense(features), float(label)) for features, label in zip(X, y)]

ops/script/lint_python.py

@@ -111,39 +111,16 @@ class LintersPaths:
         "tests/python/test_model_io.py",
         "tests/python/test_ordinal.py",
         "tests/python/test_interaction_constraints.py",
-        "tests/python-gpu/test_gpu_callbacks.py",
-        "tests/python-gpu/test_gpu_data_iterator.py",
-        "tests/python-gpu/test_gpu_ordinal.py",
-        "tests/python-gpu/load_pickle.py",
-        "tests/python-gpu/test_gpu_training_continuation.py",
-        "tests/python-gpu/test_gpu_plotting.py",
-        "tests/python-gpu/test_gpu_parse_tree.py",
+        "tests/python-gpu/",
         "tests/test_distributed/test_federated/",
         "tests/test_distributed/test_gpu_federated/",
-        "tests/test_distributed/test_with_dask/test_ranking.py",
-        "tests/test_distributed/test_with_dask/test_external_memory.py",
+        "tests/test_distributed/test_with_dask/",
         "tests/test_distributed/test_with_spark/test_data.py",
         "tests/test_distributed/test_gpu_with_spark/test_data.py",
         "tests/test_distributed/test_gpu_with_dask/",
         # demo
         "demo/dask/",
-        "demo/guide-python/custom_softmax.py",
-        "demo/guide-python/external_memory.py",
-        "demo/guide-python/distributed_extmem_basic.py",
-        "demo/guide-python/sklearn_examples.py",
-        "demo/guide-python/continuation.py",
-        "demo/guide-python/callbacks.py",
-        "demo/guide-python/update_process.py",
-        "demo/guide-python/cat_in_the_dat.py",
-        "demo/guide-python/categorical.py",
-        "demo/guide-python/cat_pipeline.py",
-        "demo/guide-python/feature_weights.py",
-        "demo/guide-python/model_parser.py",
-        "demo/guide-python/individual_trees.py",
-        "demo/guide-python/quantile_regression.py",
-        "demo/guide-python/quantile_data_iterator.py",
-        "demo/guide-python/multioutput_regression.py",
-        "demo/guide-python/learning_to_rank.py",
+        "demo/guide-python/",
         "demo/aft_survival/aft_survival_viz_demo.py",
         # CI
         "ops/",

python-package/xgboost/sklearn.py

@@ -103,7 +103,7 @@ def __call__(
         self,
         y_true: ArrayLike,
         y_pred: ArrayLike,
-        sample_weight: Optional[ArrayLike],
+        sample_weight: Optional[ArrayLike] = None,
     ) -> Tuple[ArrayLike, ArrayLike]: ...
 
 

tests/python-gpu/conftest.py

@@ -1,14 +1,16 @@
+from typing import Any, List
+
 import pytest
 
 from xgboost import testing as tm
 
 
-def has_rmm():
+def has_rmm() -> bool:
     return tm.no_rmm()["condition"]
 
 
 @pytest.fixture(scope="session", autouse=True)
-def setup_rmm_pool(request, pytestconfig):
+def setup_rmm_pool(request: Any, pytestconfig: pytest.Config) -> None:
     tm.setup_rmm_pool(request, pytestconfig)
 
 
@@ -18,7 +20,9 @@ def pytest_addoption(parser: pytest.Parser) -> None:
     )
 
 
-def pytest_collection_modifyitems(config, items):
+def pytest_collection_modifyitems(
+    config: pytest.Config, items: List[pytest.Item]
+) -> None:
     if config.getoption("--use-rmm-pool"):
         blocklist = [
             "python-gpu/test_gpu_demos.py::test_dask_training",

tests/python-gpu/test_device_quantile_dmatrix.py

@@ -203,7 +203,7 @@ def test_ref_dmatrix(self) -> None:
         strategies.fractions(0, 0.99),
     )
     @settings(print_blob=True, deadline=None)
-    def test_to_csr(self, n_samples, n_features, sparsity) -> None:
+    def test_to_csr(self, n_samples: int, n_features: int, sparsity: float) -> None:
         import cupy as cp
 
         X, y = tm.make_sparse_regression(n_samples, n_features, sparsity, False)