[Typing] bump mypy version to 1.11.2 and fix many typing errors in cinn binding (#68372)

megemini · web-flow · commit b5e30b2e0177 · 2024-09-25T18:51:00.000+08:00
diff --git a/paddle/scripts/paddle_build.sh b/paddle/scripts/paddle_build.sh
@@ -3712,24 +3712,24 @@ function exec_type_checking() {
     TITLE_CHECK_ALL=`curl -s https://github.com/PaddlePaddle/Paddle/pull/${GIT_PR_ID} | grep "<title>" | grep -i "\[typing\]" || true`
     DEBUG_MODE=`curl -s https://github.com/PaddlePaddle/Paddle/pull/${GIT_PR_ID} | grep "<title>" | grep -i "\[debug\]" || true`
 
-    # if [[ ${TITLE_CHECK_ALL} ]]; then
-    #     if [[ ${DEBUG_MODE} ]]; then
-    #         python type_checking.py --debug --full-test; type_checking_error=$?
-    #     else
-    #         python type_checking.py --full-test; type_checking_error=$?
-    #     fi
-    # else
-    #     if [[ ${DEBUG_MODE} ]]; then
-    #         python type_checking.py --debug; type_checking_error=$?
-    #     else
-    #         python type_checking.py; type_checking_error=$?
-    #     fi
-    # fi
-
-    # if [ "$type_checking_error" != "0" ];then
-    #   echo "Example code type checking failed" >&2
-    #   exit 5
-    # fi
+    if [[ ${TITLE_CHECK_ALL} ]]; then
+        if [[ ${DEBUG_MODE} ]]; then
+            python type_checking.py --debug --full-test; type_checking_error=$?
+        else
+            python type_checking.py --full-test; type_checking_error=$?
+        fi
+    else
+        if [[ ${DEBUG_MODE} ]]; then
+            python type_checking.py --debug; type_checking_error=$?
+        else
+            python type_checking.py; type_checking_error=$?
+        fi
+    fi
+
+    if [ "$type_checking_error" != "0" ];then
+      echo "Example code type checking failed" >&2
+      exit 5
+    fi
 }
 
 
diff --git a/python/paddle/base/backward.py b/python/paddle/base/backward.py
@@ -1964,7 +1964,6 @@ def _get_no_grad_set_value(no_grad_set):
 
 
 @overload
-@framework.static_only
 def append_backward(
     loss: Tensor,
     parameter_list: Sequence[Tensor | str] | None = ...,
@@ -1974,12 +1973,11 @@ def append_backward(
         | None
     ) = ...,
     checkpoints: None = ...,
-    distop_context: DistributedContext = ...,
+    distop_context: DistributedContext | None = ...,
 ) -> list[tuple[Tensor, Tensor]]: ...
 
 
 @overload
-@framework.static_only
 def append_backward(
     loss: Tensor,
     parameter_list: Sequence[Tensor | str] | None = ...,
@@ -1989,7 +1987,7 @@ def append_backward(
         | None
     ) = ...,
     checkpoints: list[Tensor] = ...,
-    distop_context: DistributedContext = ...,
+    distop_context: DistributedContext | None = ...,
 ) -> tuple[list[tuple[Tensor, Tensor]], list[str]]: ...
 
 
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -2935,7 +2935,7 @@ def size(self):
                 >>> x = paddle.static.data(name='x', shape=[3, 2, 1])
 
                 >>> # get the number of elements of the Variable
-                >>> y = x.size() # type: ignore
+                >>> y = x.size
 
         """
 
diff --git a/python/paddle/static/io.py b/python/paddle/static/io.py
@@ -892,9 +892,9 @@ def load_inference_model(
             >>> [inference_program, feed_target_names, fetch_targets] = (
             ...     paddle.static.load_inference_model(path_prefix, exe))
             >>> tensor_img = np.array(np.random.random((64, 784)), dtype=np.float32) # type: ignore[var-annotated]
-            >>> results = exe.run(inference_program,  # type: ignore[arg-type]
-            ...               feed={feed_target_names[0]: tensor_img},  # type: ignore[index,dict-item]
-            ...               fetch_list=fetch_targets)  # type: ignore[arg-type]
+            >>> results = exe.run(inference_program,
+            ...               feed={feed_target_names[0]: tensor_img},
+            ...               fetch_list=fetch_targets)
 
             # In this example, the inference program was saved in file
             # "./infer_model.pdmodel" and parameters were saved in file
diff --git a/python/paddle/utils/dlpack.py b/python/paddle/utils/dlpack.py
@@ -16,7 +16,7 @@
 
 import enum
 import warnings
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Literal, Protocol, TypeVar
 
 import paddle
 
@@ -29,18 +29,25 @@
 
     from paddle import Tensor
 
+
 __all__ = [
     'to_dlpack',
     'from_dlpack',
 ]
 
+_T_contra = TypeVar("_T_contra", contravariant=True)
+
+
+class SupportDLPack(Protocol[_T_contra]):
+    """
+    ref:
+        https://github.com/numpy/numpy/blob/7e6e48ca7aacae9994d18a3dadbabd2b91c32151/numpy/__init__.pyi#L3068-L3077
+        https://github.com/numpy/numpy/blob/7e6e48ca7aacae9994d18a3dadbabd2b91c32151/numpy/__init__.pyi#L4730-L4731
+    """
 
-class SupportDLPack(Protocol):
-    def __dlpack__(self) -> CapsuleType:
-        pass
+    def __dlpack__(self, *, stream: None | _T_contra = ...) -> CapsuleType: ...
 
-    def __dlpack_device__(self) -> tuple[enum.IntEnum, int]:
-        pass
+    def __dlpack_device__(self) -> tuple[int, Literal[0]]: ...
 
 
 class DLDeviceType(enum.IntEnum):
@@ -116,7 +123,9 @@ def to_dlpack(x: Tensor) -> CapsuleType:
     return x._to_dlpack()
 
 
-def from_dlpack(dlpack: SupportDLPack | CapsuleType) -> Tensor:
+def from_dlpack(
+    dlpack: SupportDLPack | CapsuleType,
+) -> Tensor:
     """
     Decodes a DLPack to a tensor. The returned Paddle tensor will share the memory with
     the tensor from given dlpack.
@@ -167,6 +176,7 @@ def from_dlpack(dlpack: SupportDLPack | CapsuleType) -> Tensor:
             :name: code-paddle-from-numpy
 
             >>> # Directly from external tensor that implements '__dlpack__' and '__dlpack_device__' methods
+            >>> import paddle
             >>> import numpy as np
             >>> x = np.array([[0.2, 0.3, 0.5, 0.9],
             ...              [0.1, 0.2, 0.6, 0.7]])
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -822,10 +822,6 @@ packages=['paddle',
           'paddle.decomposition',
           'paddle._typing',
           'paddle._typing.libs',
-          'paddle._typing.libs.libpaddle',
-          'paddle._typing.libs.libpaddle.pir',
-          'paddle._typing.libs.libpaddle.eager',
-          'paddle._typing.libs.libpaddle.eager.ops',
         ]
 
 with open('@PADDLE_SOURCE_DIR@/python/requirements.txt') as f:
@@ -1343,28 +1339,49 @@ if '${WITH_CPP_DIST}' == 'ON':
     paddle_lib_test_dir = '${PADDLE_LIB_TEST_DIR}'
     install_cpp_dist_and_build_test(paddle_install_dir, paddle_lib_test_dir)
 
+
 # type hints
-def extend_type_hints_package_data(package_data):
+def get_typing_libs_packages(paddle_binary_dir):
+    """get all libpaddle sub modules from 'python/paddle/_typing/libs/libpaddle'
+    e.g.
+        'paddle._typing.libs.libpaddle.cinn'
+        'paddle._typing.libs.libpaddle.pir'
+        'paddle._typing.libs.libpaddle.eager'
+        'paddle._typing.libs.libpaddle.eager.ops'
+    """
+    base_dir = Path(paddle_binary_dir) / 'python'
+    libs_dir = base_dir / 'paddle' / '_typing' / 'libs' / 'libpaddle'
+    return [
+        '.'.join(str(Path(root).relative_to(base_dir)).split(os.sep))
+        for root, _, _ in os.walk(libs_dir)
+    ]
+
+
+def extend_type_hints_package_data(packages, package_data, paddle_binary_dir):
+    typing_libs_packages = get_typing_libs_packages(paddle_binary_dir)
+
+    # update packages
+    packages += typing_libs_packages
+
+    # update package_data
     type_hints_files = {
-        'paddle': ['py.typed','*.pyi'],
+        'paddle': ['py.typed', '*.pyi'],
         'paddle.framework': ['*.pyi'],
         'paddle.base': ['*.pyi'],
         'paddle.tensor': ['tensor.pyi'],
         'paddle._typing': ['*.pyi'],
-        'paddle._typing.libs': ['*.pyi','*.md'],
-        'paddle._typing.libs.libpaddle': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.pir': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.eager': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.eager.ops': ['*.pyi'],
+        'paddle._typing.libs': ['*.pyi', '*.md'],
     }
+
+    for libpaddle_module in typing_libs_packages:
+        type_hints_files[libpaddle_module] = ['*.pyi']
+
     for pkg, files in type_hints_files.items():
         if pkg not in package_data:
             package_data[pkg] = []
         package_data[pkg] += files
 
-    return package_data
-
-package_data = extend_type_hints_package_data(package_data)
+    return packages, package_data
 
 
 def generate_stub_files(paddle_binary_dir, paddle_source_dir):
@@ -1432,6 +1449,8 @@ if os.getenv("SKIP_STUB_GEN", '').lower() not in [
 ]:
     generate_stub_files('${PADDLE_BINARY_DIR}', '${PADDLE_SOURCE_DIR}')
 
+packages, package_data = extend_type_hints_package_data(packages, package_data, '${PADDLE_BINARY_DIR}')
+
 
 with redirect_stdout():
     setup(name='${PACKAGE_NAME}',
diff --git a/python/unittest_py/requirements.txt b/python/unittest_py/requirements.txt
@@ -19,5 +19,5 @@ wandb>=0.17.2 ; python_version<"3.12"
 xlsxwriter==3.0.9
 xdoctest==1.1.1
 ubelt==1.3.3 # just for xdoctest
-mypy==1.10.0
+mypy==1.11.2
 soundfile
diff --git a/setup.py b/setup.py
@@ -1197,25 +1197,47 @@ def get_cinn_config_jsons():
     return json_path_list
 
 
-def extend_type_hints_package_data(package_data):
+def get_typing_libs_packages(paddle_binary_dir):
+    """get all libpaddle sub modules from 'python/paddle/_typing/libs/libpaddle'
+    e.g.
+        'paddle._typing.libs.libpaddle.cinn'
+        'paddle._typing.libs.libpaddle.pir'
+        'paddle._typing.libs.libpaddle.eager'
+        'paddle._typing.libs.libpaddle.eager.ops'
+    """
+    base_dir = Path(paddle_binary_dir) / 'python'
+    libs_dir = base_dir / 'paddle' / '_typing' / 'libs' / 'libpaddle'
+    return [
+        '.'.join(str(Path(root).relative_to(base_dir)).split(os.sep))
+        for root, _, _ in os.walk(libs_dir)
+    ]
+
+
+def extend_type_hints_package_data(packages, package_data, paddle_binary_dir):
+    typing_libs_packages = get_typing_libs_packages(paddle_binary_dir)
+
+    # update packages
+    packages += typing_libs_packages
+
+    # update package_data
     type_hints_files = {
         'paddle': ['py.typed', '*.pyi'],
         'paddle.framework': ['*.pyi'],
         'paddle.base': ['*.pyi'],
         'paddle.tensor': ['tensor.pyi'],
         'paddle._typing': ['*.pyi'],
         'paddle._typing.libs': ['*.pyi', '*.md'],
-        'paddle._typing.libs.libpaddle': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.pir': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.eager': ['*.pyi'],
-        'paddle._typing.libs.libpaddle.eager.ops': ['*.pyi'],
     }
+
+    for libpaddle_module in typing_libs_packages:
+        type_hints_files[libpaddle_module] = ['*.pyi']
+
     for pkg, files in type_hints_files.items():
         if pkg not in package_data:
             package_data[pkg] = []
         package_data[pkg] += files
 
-    return package_data
+    return packages, package_data
 
 
 def get_package_data_and_package_dir():
@@ -1588,9 +1610,6 @@ def get_package_data_and_package_dir():
     elif sys.platform == 'darwin':
         ext_modules = []
 
-    # type hints
-    package_data = extend_type_hints_package_data(package_data)
-
     return package_data, package_dir, ext_modules
 
 
@@ -1958,10 +1977,6 @@ def get_setup_parameters():
         'paddle.decomposition',
         'paddle._typing',
         'paddle._typing.libs',
-        'paddle._typing.libs.libpaddle',
-        'paddle._typing.libs.libpaddle.pir',
-        'paddle._typing.libs.libpaddle.eager',
-        'paddle._typing.libs.libpaddle.eager.ops',
     ]
 
     paddle_bins = ''
@@ -2235,6 +2250,10 @@ def main():
         '1',
     ]:
         generate_stub_files(paddle_binary_dir, paddle_source_dir)
+    # package stub files
+    packages, package_data = extend_type_hints_package_data(
+        packages, package_data, paddle_binary_dir
+    )
 
     setup(
         name=package_name,
diff --git a/tools/gen_pybind11_stub.py b/tools/gen_pybind11_stub.py
