[py] Use generic in xgboost model doc decorator

[MrEarle]

Context

XGBoost classes are wrapped with xboost_model_doc decorator.

Problem

xgboost_model_doc is typed as returning Callable[[Type], [Type], which ends up obfuscating the real type of the class. For example, if you try to define a union type with other classes, it ends up as Unknown:

from typing import Union
from xgboost import XGBRegressor

class MyClass1: ...
class MyClass2: ...

ModelType = XGBRegressor | MyClass1 | MyClass2
ModelType2 = Union[XGBRegressor, MyClass1, MyClass2]

Both ModelType and ModelType2 are treated as variables instead of types by mypy.

Solution

Use a generic type so that the type checker can infer that the returned class is the same as the decorated one.

With this fix, the above issue is solved.

Caveat

This is my first time contributing to this repo. I'm pretty sure this won't break anything (it's just a typing change), but please let me know if I missed anything,

[trivialfis]

Hi, thank you for the PR. I don't have any objection to the change but would like to learn more about the issue. Here's the code snippet copied from your example:

from typing import Union
from xgboost import XGBRegressor

class MyClass1: ...
class MyClass2: ...

ModelType = XGBRegressor | MyClass1 | MyClass2
ModelType2 = Union[XGBRegressor, MyClass1, MyClass2]
reveal_type(ModelType)
reveal_type(ModelType2)


def my_fn(m: ModelType) -> None:
    pass

Running mypy 1.13.0 with python 3.11.10:

$ mypy ./model_type.py 

I got the following output:

model_type.py:11: note: Revealed type is "types.UnionType[xgboost.sklearn.XGBRegressor, model_type.MyClass1, model_type.MyClass2]"
model_type.py:12: note: Revealed type is "typing._SpecialForm"
Success: no issues found in 1 source file

It doesn't contain any errors. I'm using the master branch of XGBoost. Also, have you tried the typing.TypeAlias?

[MrEarle]

Huh now that you mention it, I only get a mypy error when combining it with a sklearn model:

from xgboost import XGBRegressor
from sklearn.ensemble import RandomForestRegressor # type: ignore

class MyClass:
    def fit(self, X, y):
        return self

_ModelType = MyClass | XGBRegressor | RandomForestRegressor

reveal_type(_ModelType)

def train_model(model: _ModelType) -> _ModelType:
    return model.fit([[1]], [1])

And mypy's output:

But if I remove sklearn's model, mypy doesn't complain, so part of the issue is probably with sklearn as well (my bad).

But even without the union type, vscode has a hard time infering the type:

Without the fix:

With the fix:

So I guess the issue is really with VSCode.

[trivialfis]

I don't use vscode, if not mistaken, it might be using https://github.com/microsoft/pyright instead of mypy? I have tried pyright before, but the result was not consistent with mypy, not sure about the current status.

Anyway, we can merge the PR if it helps with vscode users and the tests can pass.

[MrEarle]

Yeah, sorry. Forgot to run black. Should be fine now.