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

[python-package] add more type hints on Booster (#5360)



* [python-package] add more type hints on Booster

* add return to Callable

* Apply suggestions from code review
Co-authored-by: default avatarNikita Titov <nekit94-08@mail.ru>
Co-authored-by: default avatarNikita Titov <nekit94-08@mail.ru>
parent 58bc3af7
......@@ -3380,7 +3380,11 @@ class Booster:
_dump_pandas_categorical(self.pandas_categorical, filename)
return self
def shuffle_models(self, start_iteration=0, end_iteration=-1):
def shuffle_models(
self,
start_iteration: int = 0,
end_iteration: int = -1
) -> "Booster":
"""Shuffle models.
Parameters
......@@ -3432,7 +3436,12 @@ class Booster:
self.pandas_categorical = _load_pandas_categorical(model_str=model_str)
return self
def model_to_string(self, num_iteration=None, start_iteration=0, importance_type='split'):
def model_to_string(
self,
num_iteration: Optional[int] = None,
start_iteration: int = 0,
importance_type: str = 'split'
) -> str:
"""Save Booster to string.
Parameters
......@@ -3485,7 +3494,13 @@ class Booster:
ret += _dump_pandas_categorical(self.pandas_categorical)
return ret
def dump_model(self, num_iteration=None, start_iteration=0, importance_type='split', object_hook=None):
def dump_model(
self,
num_iteration: Optional[int] = None,
start_iteration: int = 0,
importance_type: str = 'split',
object_hook: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None
) -> Dict[str, Any]:
"""Dump Booster to JSON format.
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