[python-package] fix mypy errors for eval_at and basic (#5674)

microsoft · Jan 13, 2023 · 5df7d51 · 5df7d51
1 parent 02d212b
commit 5df7d51
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
@@ -2186,7 +2186,7 @@ def get_field(self, field_name: str) -> Optional[np.ndarray]:
 
     def set_categorical_feature(
         self,
-        categorical_feature: Union[List[int], List[str]]
+        categorical_feature: Union[List[int], List[str], str]
     ) -> "Dataset":
         """Set categorical features.
 
@@ -2268,7 +2268,7 @@ def set_reference(self, reference: "Dataset") -> "Dataset":
             raise LightGBMError("Cannot set reference after freed raw data, "
                                 "set free_raw_data=False when construct Dataset to avoid this.")
 
-    def set_feature_name(self, feature_name: List[str]) -> "Dataset":
+    def set_feature_name(self, feature_name: Union[List[str], str]) -> "Dataset":
         """Set feature name.
 
         Parameters

diff --git a/python-package/lightgbm/dask.py b/python-package/lightgbm/dask.py
@@ -413,7 +413,7 @@ def _train(
     eval_init_score: Optional[List[_DaskCollection]] = None,
     eval_group: Optional[List[_DaskVectorLike]] = None,
     eval_metric: Optional[_LGBM_ScikitEvalMetricType] = None,
-    eval_at: Optional[Union[List[int], Tuple[int]]] = None,
+    eval_at: Optional[Union[List[int], Tuple[int, ...]]] = None,
     **kwargs: Any
 ) -> LGBMModel:
     """Inner train routine.
@@ -1045,7 +1045,7 @@ def _lgb_dask_fit(
         eval_init_score: Optional[List[_DaskCollection]] = None,
         eval_group: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[_LGBM_ScikitEvalMetricType] = None,
-        eval_at: Optional[Iterable[int]] = None,
+        eval_at: Optional[Union[List[int], Tuple[int, ...]]] = None,
         **kwargs: Any
     ) -> "_DaskLGBMModel":
         if not DASK_INSTALLED:
@@ -1499,7 +1499,7 @@ def fit(
         eval_init_score: Optional[List[_DaskVectorLike]] = None,
         eval_group: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[_LGBM_ScikitEvalMetricType] = None,
-        eval_at: Union[List[int], Tuple[int]] = (1, 2, 3, 4, 5),
+        eval_at: Union[List[int], Tuple[int, ...]] = (1, 2, 3, 4, 5),
         **kwargs: Any
     ) -> "DaskLGBMRanker":
         """Docstring is inherited from the lightgbm.LGBMRanker.fit."""

diff --git a/python-package/lightgbm/sklearn.py b/python-package/lightgbm/sklearn.py
@@ -1204,7 +1204,7 @@ def fit(
         eval_init_score=None,
         eval_group=None,
         eval_metric: Optional[_LGBM_ScikitEvalMetricType] = None,
-        eval_at: Union[List[int], Tuple[int]] = (1, 2, 3, 4, 5),
+        eval_at: Union[List[int], Tuple[int, ...]] = (1, 2, 3, 4, 5),
         feature_name='auto',
         categorical_feature='auto',
         callbacks=None,