Repository: Lightning-AI/pytorch-lightning
Pull request: load_from_checkpoint returns the expected type
## What does this PR do?

- `Self`/`self` cannot be used since it requires python 3.11 AFAIK
- Since it looks like the tests are not type checked, I added a folder under `tests` for type checking (maybe more will be added)

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #15467

### Does your PR introduce any breaking changes? If yes, please list them.

None

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [ ] Did you list all the **breaking changes** introduced by this pull request?
- [ ] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
`MyNet.load_from_checkpoint` should return an instance of type `MyNet`
## Proposed refactor

`MyNet.load_from_checkpoint` currently returns an instance of type `LightningModule`.
It would be nice if it would return an instance of type `MyNet`.

### Motivation

<!-- Please outline the motivation for the proposal. If this is related to another GitHub issue, please link it here -->

This seems to make sense from an API point of view, and helps IDE's and static type checkers with code accessing properties/methods of type `MyNet`. 

### Pitch
```
class MyNet(pl.LightningModule):
  
    @property
    def num_classes(self):
        return self._model.out_channels

mynet = MyNet.load_from_checkpoint(file_path)
if mynet.num_classes > 2:
  print("hello world")
```

mypy should not complain about code like the (pseudo-code) above. 

### Additional context

<!-- Add any other context or screenshots here. -->


cc @borda @justusschock @awaelchli @rohitgr7

Hints from issue discussion:
@dyollb This can be achieved with typing.TypeVar. Here is a quick sketch:

```py
from typing import TypeVar

LM = TypeVar("LM")


@classmethod
def load_from_checkpoint(cls: LM, ...) -> LM:
```

Would you be interested to investigate and contribute this? 
I had a quick try but didn't figure out how to do it, because load_from_checkpoint is defined in the mixin class ModelIO.
Thanks for trying!! 
Let's see if someone else from the community wants to give this a try. If nobody takes it, I will take a look too <3 
This can be done by just using the `Self` type from https://peps.python.org/pep-0673/

Commit message: attempt to fix mypy errors for py37

Patch to review:
diff --git a/src/pytorch_lightning/core/datamodule.py b/src/pytorch_lightning/core/datamodule.py
index 5cd1713e2ffba..2e2bd7cbd506d 100644
--- a/src/pytorch_lightning/core/datamodule.py
+++ b/src/pytorch_lightning/core/datamodule.py
@@ -17,6 +17,7 @@
 from typing import Any, Dict, IO, List, Mapping, Optional, Sequence, Tuple, Union
 
 from torch.utils.data import DataLoader, Dataset, IterableDataset
+from typing_extensions import Self
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.types import _PATH
@@ -218,7 +219,7 @@ def load_from_checkpoint(
         checkpoint_path: Union[_PATH, IO],
         hparams_file: Optional[_PATH] = None,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
         r"""
         Primary way of loading a datamodule from a checkpoint. When Lightning saves a checkpoint
         it stores the arguments passed to ``__init__``  in the checkpoint under ``"datamodule_hyper_parameters"``.
diff --git a/src/pytorch_lightning/core/saving.py b/src/pytorch_lightning/core/saving.py
index fefdb122fc5d2..df9a7d2bf7507 100644
--- a/src/pytorch_lightning/core/saving.py
+++ b/src/pytorch_lightning/core/saving.py
@@ -26,6 +26,7 @@
 
 import yaml
 from lightning_utilities.core.apply_func import apply_to_collection
+from typing_extensions import Self
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.cloud_io import _load as pl_load
@@ -63,7 +64,7 @@ def load_from_checkpoint(
         hparams_file: Optional[str] = None,
         strict: bool = True,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
         r"""
         Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint
         it stores the arguments passed to ``__init__``  in the checkpoint under ``"hyper_parameters"``.


Code review comments:
diff_hunk: @@ -63,7 +64,7 @@ def load_from_checkpoint(
         hparams_file: Optional[str] = None,
         strict: bool = True,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
path: src/pytorch_lightning/core/saving.py
text: @user1:
```suggestion
    ) -> Self:  # type: ignore[valid-type]
```

diff_hunk: @@ -218,7 +219,7 @@ def load_from_checkpoint(
         checkpoint_path: Union[_PATH, IO],
         hparams_file: Optional[_PATH] = None,
         **kwargs: Any,
-    ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
+    ) -> Self:  # type: ignore [valid-type]
path: src/pytorch_lightning/core/datamodule.py
text: @user1:
```suggestion
    ) -> Self:  # type: ignore[valid-type]
```