Fix mypy.

dmlc · Jan 8, 2021 · f406a48 · f406a48
1 parent 5b6005c
commit f406a48
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
@@ -1325,7 +1325,7 @@ async def _fit_async(
  sample_weight: Optional[_DaskCollection],
  base_margin: Optional[_DaskCollection],
  eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
- eval_metric: Optional[Union[str, List[str], Callable]],
+ eval_metric: Optional[Union[str, List[str], Metric]],
  sample_weight_eval_set: Optional[List[_DaskCollection]],
  early_stopping_rounds: int,
  verbose: bool,
@@ -1382,7 +1382,7 @@ def fit(
  sample_weight: Optional[_DaskCollection] = None,
  base_margin: Optional[_DaskCollection] = None,
  eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]] = None,
- eval_metric: Optional[Callable] = None,
+ eval_metric: Optional[Union[str, List[str], Metric]] = None,
  early_stopping_rounds: Optional[int] = None,
  verbose: bool = True,
  xgb_model: Optional[Union[Booster, XGBModel]] = None,
@@ -1418,7 +1418,7 @@ async def _fit_async(
  sample_weight: Optional[_DaskCollection],
  base_margin: Optional[_DaskCollection],
  eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
- eval_metric: Optional[Union[str, List[str], Callable]],
+ eval_metric: Optional[Union[str, List[str], Metric]],
  sample_weight_eval_set: Optional[List[_DaskCollection]],
  early_stopping_rounds: int,
  verbose: bool,
@@ -1493,7 +1493,7 @@ def fit(
  sample_weight: Optional[_DaskCollection] = None,
  base_margin: Optional[_DaskCollection] = None,
  eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]] = None,
- eval_metric: Optional[Callable] = None,
+ eval_metric: Optional[Union[str, List[str], Metric]] = None,
  early_stopping_rounds: Optional[int] = None,
  verbose: bool = True,
  xgb_model: Optional[Union[Booster, XGBModel]] = None,
@@ -1601,7 +1601,7 @@ async def _fit_async(
  eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
  sample_weight_eval_set: Optional[List[_DaskCollection]],
  eval_qid: Optional[List[_DaskCollection]],
- eval_metric: Optional[Union[str, List[str], Callable]],
+ eval_metric: Optional[Union[str, List[str], Metric]],
  early_stopping_rounds: int,
  verbose: bool,
  xgb_model: Optional[Union[XGBModel, Booster]],
@@ -1665,7 +1665,7 @@ def fit( # pylint: disable=arguments-differ
  sample_weight_eval_set: Optional[List[_DaskCollection]] = None,
  eval_group: Optional[List[_DaskCollection]] = None,
  eval_qid: Optional[List[_DaskCollection]] = None,
- eval_metric: Optional[Union[str, List[str], Callable]] = None,
+ eval_metric: Optional[Union[str, List[str], Metric]] = None,
  early_stopping_rounds: int = None,
  verbose: bool = False,
  xgb_model: Optional[Union[XGBModel, Booster]] = None,

diff --git a/python-package/xgboost/sklearn.py b/python-package/xgboost/sklearn.py
@@ -7,6 +7,7 @@
 from typing import Union, Optional, List, Dict, Callable, Tuple, Any
 import numpy as np
 from .core import Booster, DMatrix, XGBoostError, _deprecate_positional_args
+from .core import Metric
 from .training import train
 from .data import _is_cudf_df, _is_cudf_ser, _is_cupy_array
 
@@ -558,7 +559,7 @@ def _configure_fit(
  booster: Optional[Booster],
  eval_metric: Optional[Union[Callable, str, List[str]]],
  params: Dict[str, Any],
- ) -> Tuple[Booster, Optional[Union[Callable, str, List[str]]], Dict[str, Any]]:
+ ) -> Tuple[Booster, Optional[Metric], Dict[str, Any]]:
  # pylint: disable=protected-access, no-self-use
  model = booster
  if hasattr(model, '_Booster'):