Unverified Commit 52a50dde authored by James Lamb's avatar James Lamb Committed by GitHub
Browse files

[python-package] add more type hints in basic.py (#5407)

* [python-package] add more type hints in basic.py

* Booster handle is a c_void_p
parent 0a5c5838
...@@ -722,7 +722,12 @@ class _InnerPredictor: ...@@ -722,7 +722,12 @@ class _InnerPredictor:
Can be converted from Booster, but cannot be converted to Booster. Can be converted from Booster, but cannot be converted to Booster.
""" """
def __init__(self, model_file=None, booster_handle=None, pred_parameter=None): def __init__(
self,
model_file: Optional[Union[str, Path]] = None,
booster_handle: Optional[ctypes.c_void_p] = None,
pred_parameter: Optional[Dict[str, Any]] = None
):
"""Initialize the _InnerPredictor. """Initialize the _InnerPredictor.
Parameters Parameters
...@@ -773,7 +778,7 @@ class _InnerPredictor: ...@@ -773,7 +778,7 @@ class _InnerPredictor:
except AttributeError: except AttributeError:
pass pass
def __getstate__(self): def __getstate__(self) -> Dict[str, Any]:
this = self.__dict__.copy() this = self.__dict__.copy()
this.pop('handle', None) this.pop('handle', None)
return this return this
...@@ -3792,7 +3797,10 @@ class Booster: ...@@ -3792,7 +3797,10 @@ class Booster:
ctypes.byref(ret))) ctypes.byref(ret)))
return ret.value return ret.value
def _to_predictor(self, pred_parameter=None): def _to_predictor(
self,
pred_parameter: Optional[Dict[str, Any]] = None
) -> _InnerPredictor:
"""Convert to predictor.""" """Convert to predictor."""
predictor = _InnerPredictor(booster_handle=self.handle, pred_parameter=pred_parameter) predictor = _InnerPredictor(booster_handle=self.handle, pred_parameter=pred_parameter)
predictor.pandas_categorical = self.pandas_categorical predictor.pandas_categorical = self.pandas_categorical
...@@ -3887,7 +3895,12 @@ class Booster: ...@@ -3887,7 +3895,12 @@ class Booster:
else: else:
return result return result
def get_split_value_histogram(self, feature, bins=None, xgboost_style=False): def get_split_value_histogram(
self,
feature: Union[int, str],
bins: Optional[Union[int, str]] = None,
xgboost_style: bool = False
) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray, pd_DataFrame]:
"""Get split value histogram for the specified feature. """Get split value histogram for the specified feature.
Parameters Parameters
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment