[Typing] Upgrade mypy to 2.1.0 and target Python 3.10

pyproject.toml

@@ -172,8 +172,9 @@ known-first-party = ["paddle"]
 "test/dygraph_to_static/seq2seq_dygraph_model.py" = ["RUF005"]
 
 [tool.mypy]
-python_version = "3.9"
+python_version = "3.10"
 cache_dir = ".mypy_cache"
+num_workers = 8
 # Miscellaneous strictness flags
 allow_redefinition = true
 local_partial_types = true

python/paddle/base/framework.py

@@ -8350,7 +8350,7 @@ def device_guard(device: str | None = None) -> Generator[None, None, None]:
             >>> support_gpu = paddle.is_compiled_with_cuda()
             >>> place = paddle.CPUPlace()
             >>> if support_gpu:
-            ...     place = paddle.CUDAPlace(0)  # type: ignore
+            ...     place = paddle.CUDAPlace(0)
 
             >>> # if GPU is supported, the three OPs below will be automatically assigned to CUDAPlace(0)
             >>> data1 = paddle.full(shape=[1, 3, 8, 8], fill_value=0.5, dtype='float32')

python/paddle/distributed/communication/recv.py

@@ -142,7 +142,7 @@ def recv_object_list(
             ...     data = ["hello", {"key": 100}, [1, 2, 3]]
             ...     dist.send_object_list(data, dst=1)
             >>> else:
-            ...     data = [None] * 3  # type: ignore
+            ...     data = [None] * 3
             ...     dist.recv_object_list(data, src=0)
             >>> print(data)
             >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)

python/paddle/distributed/communication/send.py

@@ -141,7 +141,7 @@ def send_object_list(
             ...     data = ["hello", {"key": 100}, [1, 2, 3]]
             ...     dist.send_object_list(data, dst=1)
             >>> else:
-            ...     data = [None] * 3  # type: ignore
+            ...     data = [None] * 3
             ...     dist.recv_object_list(data, src=0)
             >>> print(data)
             >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)

python/paddle/distributed/fleet/base/topology.py

@@ -93,7 +93,7 @@ def create_nccl_config(
             >>> import paddle.distributed as dist
             >>> from typing import Union
             >>> dist.init_parallel_env()
-            >>> nccl_config: dict[str, Union[int, str]] = {
+            >>> raw_nccl_config: dict[str, Union[int, str]] = {
             ...     "commName": "tp_comm",
             ...     "ll_buffsize": 0,
             ...     "ll128_buffsize": 0,
@@ -104,7 +104,7 @@ def create_nccl_config(
             ...     "protoStr": "Simple",
             ... }
             >>> ranks = [0, 1, 2, 3, 4, 5, 6, 7]
-            >>> nccl_config = dist.create_nccl_config(nccl_config)
+            >>> nccl_config = dist.create_nccl_config(raw_nccl_config)
             >>> pg = dist.new_group(ranks, nccl_config=nccl_config)
             >>> m, n = 4096, 8192
             >>> local_rank = dist.get_rank(pg)

python/paddle/nn/functional/loss.py

@@ -3143,7 +3143,7 @@ def cross_entropy(
             >>> shape = [N, C]
             >>> label_smoothing = 0.4
             >>> reduction = 'mean'
-            >>> weight: Optional[paddle.Tensor] = None
+            >>> weight = None
             >>> logits = paddle.uniform(shape, dtype='float64', min=0.1, max=1.0)
             >>> integer_labels = paddle.randint(low=0, high=C, size=[N], dtype='int64')
             >>> one_hot_labels = paddle.nn.functional.one_hot(integer_labels, C).astype('float32')

python/paddle/nn/initializer/assign.py

@@ -33,6 +33,8 @@
 
     import numpy.typing as npt
 
+    from paddle._typing import NestedSequence
+
 __all__ = []
 
 
@@ -271,7 +273,9 @@ class Assign(NumpyArrayInitializer):
 
     def __init__(
         self,
-        value: npt.NDArray[Any] | Sequence[int] | paddle.Tensor,
+        value: npt.NDArray[Any]
+        | Sequence[NestedSequence[int | float | bool | complex]]
+        | paddle.Tensor,
         name: str | None = None,
     ) -> None:
         import numpy

python/paddle/nn/layer/loss.py

@@ -405,7 +405,7 @@ class CrossEntropyLoss(Layer):
             >>> shape = [N, C]
             >>> label_smoothing = 0.4
             >>> reduction = 'mean'
-            >>> weight: Optional[paddle.Tensor] = None
+            >>> weight = None
             >>> logits = paddle.uniform(shape, dtype='float64', min=0.1, max=1.0)
             >>> integer_labels = paddle.randint(low=0, high=C, size=[N], dtype='int64')
             >>> one_hot_labels = paddle.nn.functional.one_hot(integer_labels, C).astype('float32')

python/paddle/vision/datasets/folder.py

@@ -207,9 +207,14 @@ class DatasetFolder(Dataset[tuple["_ImageDataType", int]]):
             ...     ]
             ... )
 
+            >>> def cv2_loader(path: str):
+            ...     image = cv2.imread(path)
+            ...     assert image is not None
+            ...     return image
+
             >>> data_folder_2 = DatasetFolder(
             ...     fake_data_dir,
-            ...     loader=lambda x: cv2.imread(x),  # load image with OpenCV
+            ...     loader=cv2_loader,  # load image with OpenCV
             ...     extensions=(".jpg",),  # only load *.jpg files
             ...     transform=transform,  # apply transform to every image
             ... )
@@ -442,9 +447,14 @@ class ImageFolder(Dataset[list["_ImageDataType"]]):
             ...     ]
             ... )
 
+            >>> def cv2_loader(path: str):
+            ...     image = cv2.imread(path)
+            ...     assert image is not None
+            ...     return image
+
             >>> image_folder_2 = ImageFolder(
             ...     fake_data_dir,
-            ...     loader=lambda x: cv2.imread(x),  # load image with OpenCV
+            ...     loader=cv2_loader,  # load image with OpenCV
             ...     extensions=(".jpg",),  # only load *.jpg files
             ...     transform=transform,  # apply transform to every image
             ... )

python/paddle/vision/transforms/functional.py

@@ -437,12 +437,16 @@ def adjust_brightness(
             >>> fake_img = Image.fromarray(fake_img)
             >>> print(fake_img.size)
             (300, 256)
-            >>> print(fake_img.load()[1, 1])  # type: ignore[index]
+            >>> fake_img_pixels = fake_img.load()
+            >>> assert fake_img_pixels is not None
+            >>> print(fake_img_pixels[1, 1])
             (61, 155, 171)
             >>> converted_img = F.adjust_brightness(fake_img, 0.5)
             >>> print(converted_img.size)
             (300, 256)
-            >>> print(converted_img.load()[1, 1])
+            >>> converted_img_pixels = converted_img.load()
+            >>> assert converted_img_pixels is not None
+            >>> print(converted_img_pixels[1, 1])
             (30, 77, 85)
     """
     if not (

python/paddle/vision/transforms/transforms.py

@@ -24,12 +24,11 @@
     Generic,
     Literal,
     Protocol,
-    TypeVar,
     overload,
 )
 
 import numpy as np
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, TypeVar
 
 import paddle
 
@@ -45,6 +44,7 @@
     _TransformInputKeys: TypeAlias = Sequence[
         Literal["image", "coords", "boxes", "mask"]
     ]
+    _ImageDataType: TypeAlias = Tensor | PILImage | npt.NDArray[Any]
     from .functional import (
         _InterpolationCv2,
         _InterpolationPil,
@@ -53,11 +53,9 @@
 
 
 _InputT = TypeVar(
-    "_InputT", "Tensor", "PILImage", "npt.NDArray[Any]", contravariant=True
-)
-_RetT = TypeVar(
-    "_RetT", "Tensor", "PILImage", "npt.NDArray[Any]", covariant=True
+    "_InputT", bound="_ImageDataType", contravariant=True, default=Any
 )
+_RetT = TypeVar("_RetT", bound="_ImageDataType", covariant=True, default=Any)
 
 
 class _Transform(Protocol, Generic[_InputT, _RetT]):
@@ -226,7 +224,7 @@ class BaseTransform(_Transform[_InputT, _RetT]):
             ...         return img.shape[:2][::-1]
             ...     else:
             ...         raise TypeError("Unexpected type {}".format(type(img)))
-            >>> class CustomRandomFlip(BaseTransform):  # type: ignore[type-arg]
+            >>> class CustomRandomFlip(BaseTransform):
             ...     def __init__(self, prob=0.5, keys=None):
             ...         super().__init__(keys)
             ...         self.prob = prob
@@ -1038,11 +1036,15 @@ class BrightnessTransform(BaseTransform[_InputT, _RetT]):
 
             >>> transform = BrightnessTransform(0.4)
             >>> fake_img = Image.fromarray((np.random.rand(224, 224, 3) * 255.0).astype(np.uint8))
-            >>> print(fake_img.load()[1, 1])  # type: ignore[index]
+            >>> fake_img_pixels = fake_img.load()
+            >>> assert fake_img_pixels is not None
+            >>> print(fake_img_pixels[1, 1])
             (60, 169, 34)
             >>> # doctest: +SKIP('random sample in Brightness function')
             >>> fake_img = transform(fake_img)
-            >>> print(fake_img.load()[1, 1])
+            >>> converted_img_pixels = fake_img.load()
+            >>> assert converted_img_pixels is not None
+            >>> print(converted_img_pixels[1, 1])
             (68, 192, 38)
 
     """

python/unittest_py/requirements.txt

@@ -17,7 +17,7 @@ wandb>=0.23.1
 xlsxwriter==3.0.9
 xdoctest==1.3.0
 ubelt==1.4.0 # just for xdoctest
-mypy==1.19.1
+mypy==2.1.0
 soundfile
 apache-tvm-ffi==0.1.11
 graphviz

tools/type_checking.py

@@ -101,7 +101,9 @@ class TestResult:
     extra_info: dict[str, Any] = field(default_factory=dict)
 
 
-def pty_run(command: list[str]) -> subprocess.CompletedProcess[str]:
+def pty_run(
+    command: list[str], env: dict[str, str] | None = None
+) -> subprocess.CompletedProcess[str]:
     """Run a command in a pseudo-terminal to capture colored output."""
     master_fd, slave_fd = pty.openpty()
     try:
@@ -113,6 +115,7 @@ def pty_run(command: list[str]) -> subprocess.CompletedProcess[str]:
             stdout=slave_fd,
             stderr=slave_fd,
             close_fds=True,
+            env=env,
         )
 
         # Parent no longer needs the slave fd — close it so the child can
@@ -206,6 +209,14 @@ def _parse_output(
         assert summary, 'No summary found in mypy output'
         return results, summary
 
+    def _mypy_env(self) -> dict[str, str]:
+        env = os.environ.copy()
+        if 'MYPY_NUM_WORKERS' not in env:
+            num_workers = os.cpu_count()
+            if num_workers is not None:
+                env['MYPY_NUM_WORKERS'] = str(num_workers)
+        return env
+
     def run_on_directory(
         self,
         dir: pathlib.Path,
@@ -221,6 +232,7 @@ def run_on_directory(
                 "--pretty",
                 str(dir),
             ],
+            env=self._mypy_env(),
         )
         if res.returncode == 0:
             print(f'No type errors found in directory {dir}')