Incorrect type hints for LightningModule.load_from_checkpoint

[adamjstewart]

Bug description

The LightningModule.load_from_checkpoint function appears to return type Self? instead of LightningModule.

How to reproduce the bug

First, create a test.py file containing the following code:

import pytorch_lightning as pl

model1 = pl.LightningModule()
reveal_type(model1)
print(model1.hparams)

model2 = pl.LightningModule.load_from_checkpoint("/foo/bar")
reveal_type(model2)
print(model2.hparams)

Then, run mypy on this file:

$ mypy test.py
test.py:5: note: Revealed type is "pytorch_lightning.core.module.LightningModule"
test.py:9: note: Revealed type is "Self?"
test.py:10: error: Self? has no attribute "hparams"  [attr-defined]
Found 1 error in 1 file (checked 1 source file)

You'll notice that the type of the second model is incorrect, so mypy raises an error when you try to access LightningModule attributes like hparams.

Error messages and logs

We hit this issue in CI after the release of 1.8.1, so this may have been introduced in 1.8.1.

https://github.com/microsoft/torchgeo/actions/runs/3499809303/jobs/5861774881

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
#- PyTorch Lightning Version (e.g., 1.5.0): 1.8.2
#- Lightning App Version (e.g., 0.5.2): N/A
#- PyTorch Version (e.g., 1.10): 1.12.1
#- Python version (e.g., 3.9): 3.10.8
#- OS (e.g., Linux): macOS and Linux
#- CUDA/cuDNN version: N/A
#- GPU models and configuration: N/A
#- How you installed Lightning(`conda`, `pip`, source): Spack
#- Running environment of LightningApp (e.g. local, cloud): local

More info

I'm happy to submit a PR to fix this issue if someone wants to point me in the right direction. I have dozens of other typing issues with PL, so would love to be able to fix those as well.

cc @Borda

[adamjstewart]

I'm able to reproduce this outside of PL, so I opened an issue with typing-extensions: python/typing_extensions#96

[adamjstewart]

Feedback from python/typing_extensions#96 is that the latest release of mypy does not yet know how to handle Self. Support was added in python/mypy#14041 which should be in the next release. So we either need to manually cast to the right type and wait for the next mypy release, or use TypeVar instead.

[awaelchli]

@adamjstewart This was recently added in #15496. See the explanation here: #15496 (comment)

cc @carmocca

[carmocca]

Exactly. You can install the latest mypy by following the instructions at https://github.com/python/mypy#quick-start. I expect they'll make a release within the next month.

[adamjstewart]

Decided to manually cast things to the right type for now. Will look out for the new mypy release. Thanks for the pointers!