[Typing][B-53] Add type annotations for python/paddle/utils/cpp_extension/cpp_extension.py (#65818)

megemini · web-flow · commit 5b1b4060dc2a · 2024-07-09T18:26:29.000+08:00
diff --git a/python/paddle/utils/cpp_extension/cpp_extension.py b/python/paddle/utils/cpp_extension/cpp_extension.py
@@ -14,6 +14,10 @@
 
 # isort: skip_file
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
 import os
 import copy
 import re
@@ -61,6 +65,9 @@
 
 from ...base import core
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
 # Note(zhouwei): On windows, it will export function 'PyInit_[name]' by default,
 # The solution is: 1.User add function PyInit_[name] 2. set not to export
 # refer to https://stackoverflow.com/questions/34689210/error-exporting-symbol-when-building-python-c-extension-in-windows
@@ -76,7 +83,7 @@
     CUDA_HOME = ROCM_HOME
 
 
-def setup(**attr):
+def setup(**attr: Any) -> None:
     """
     The interface is used to config the process of compiling customized operators,
     mainly includes how to compile shared library, automatically generate python API
@@ -234,7 +241,9 @@ def setup(**attr):
         setuptools.setup(**attr)
 
 
-def CppExtension(sources, *args, **kwargs):
+def CppExtension(
+    sources: Sequence[str], *args: Any, **kwargs: Any
+) -> setuptools.Extension:
     """
     The interface is used to config source files of customized operators and complies
     Op Kernel only supporting CPU device. Please use ``CUDAExtension`` if you want to
@@ -284,7 +293,9 @@ def CppExtension(sources, *args, **kwargs):
     return setuptools.Extension(name, sources, *args, **kwargs)
 
 
-def CUDAExtension(sources, *args, **kwargs):
+def CUDAExtension(
+    sources: Sequence[str], *args: Any, **kwargs: Any
+) -> setuptools.Extension:
     """
     The interface is used to config source files of customized operators and complies
     Op Kernel supporting both CPU and GPU devices. Please use ``CppExtension`` if you want to
@@ -359,7 +370,7 @@ class BuildExtension(build_ext):
     """
 
     @classmethod
-    def with_options(cls, **options):
+    def with_options(cls, **options: Any) -> type[BuildExtension]:
         """
         Returns a BuildExtension subclass containing use-defined options.
         """
@@ -371,7 +382,7 @@ def __init__(self, *args, **kwargs):
 
         return cls_with_options
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """
         Attributes is initialized with following order:
 
@@ -388,18 +399,18 @@ def __init__(self, *args, **kwargs):
         # whether containing cuda source file in Extensions
         self.contain_cuda_file = False
 
-    def initialize_options(self):
+    def initialize_options(self) -> None:
         super().initialize_options()
 
-    def finalize_options(self):
+    def finalize_options(self) -> None:
         super().finalize_options()
         # NOTE(Aurelius84): Set location of compiled shared library.
         # Carefully to modify this because `setup.py build/install`
         # and `load` interface rely on this attribute.
         if self.output_dir is not None:
             self.build_lib = self.output_dir
 
-    def build_extensions(self):
+    def build_extensions(self) -> None:
         if OS_NAME.startswith("darwin"):
             self._valid_clang_compiler()
 
@@ -640,7 +651,7 @@ def wrapper(source_filenames, strip_dir=0, output_dir=''):
         so_path = self.get_ext_fullpath(self.extensions[0]._full_name)
         _reset_so_rpath(so_path)
 
-    def get_ext_filename(self, fullname):
+    def get_ext_filename(self, fullname: str) -> str:
         # for example: customized_extension.cpython-37m-x86_64-linux-gnu.so
         ext_name = super().get_ext_filename(fullname)
         split_str = '.'
@@ -658,7 +669,7 @@ def get_ext_filename(self, fullname):
             ext_name = split_str.join(name_items)
         return ext_name
 
-    def _valid_clang_compiler(self):
+    def _valid_clang_compiler(self) -> None:
         """
         Make sure to use Clang as compiler on Mac platform
         """
@@ -672,7 +683,7 @@ def _valid_clang_compiler(self):
             linker_so=linker_infos,
         )
 
-    def _check_abi(self):
+    def _check_abi(self) -> None:
         """
         Check ABI Compatibility.
         """
@@ -697,7 +708,7 @@ def _check_abi(self):
             )
             raise UserWarning(msg)
 
-    def _record_op_info(self):
+    def _record_op_info(self) -> None:
         """
         Record custom op information.
         """
@@ -729,11 +740,11 @@ class EasyInstallCommand(easy_install):
                     library file after extracting egg-info into site-packages.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
     # NOTE(Aurelius84): Add args and kwargs to make compatible with PY2/PY3
-    def run(self, *args, **kwargs):
+    def run(self, *args: Any, **kwargs: Any) -> None:
         super().run(*args, **kwargs)
         # NOTE: To avoid failing import .so file instead of
         # python file because they have same name, we rename
@@ -764,7 +775,7 @@ class BuildCommand(build):
     """
 
     @classmethod
-    def with_options(cls, **options):
+    def with_options(cls, **options: Any) -> type[BuildCommand]:
         """
         Returns a BuildCommand subclass containing use-defined options.
         """
@@ -776,13 +787,13 @@ def __init__(self, *args, **kwargs):
 
         return cls_with_options
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Note: shall put before super()
         self._specified_build_base = kwargs.get('build_base', None)
 
         super().__init__(*args, **kwargs)
 
-    def initialize_options(self):
+    def initialize_options(self) -> None:
         """
         build_base is root directory for all sub-command, such as
         build_lib, build_temp. See `distutils.command.build` for details.
@@ -793,16 +804,16 @@ def initialize_options(self):
 
 
 def load(
-    name,
-    sources,
-    extra_cxx_cflags=None,
-    extra_cuda_cflags=None,
-    extra_ldflags=None,
-    extra_include_paths=None,
-    extra_library_paths=None,
-    build_directory=None,
-    verbose=False,
-):
+    name: str,
+    sources: Sequence[str],
+    extra_cxx_cflags: Sequence[str] | None = None,
+    extra_cuda_cflags: Sequence[str] | None = None,
+    extra_ldflags: Sequence[str] | None = None,
+    extra_include_paths: Sequence[str] | None = None,
+    extra_library_paths: Sequence[str] | None = None,
+    build_directory: str | None = None,
+    verbose: bool = False,
+) -> ModuleType:
     """
     An Interface to automatically compile C++/CUDA source files Just-In-Time
     and return callable python function as other Paddle layers API. It will
@@ -863,22 +874,22 @@ def load(
         name(str): Specify the name of generated shared library file name, not including ``.so`` and ``.dll`` suffix.
         sources(list[str]): Specify source files name of customized operators.  Supporting ``.cc`` , ``.cpp`` for CPP file
                             and ``.cu`` for CUDA file.
-        extra_cxx_cflags(list[str], optional): Specify additional flags used to compile CPP files. By default
+        extra_cxx_cflags(list[str]|None, optional): Specify additional flags used to compile CPP files. By default
                                all basic and framework related flags have been included.
-        extra_cuda_cflags(list[str], optional): Specify additional flags used to compile CUDA files. By default
+        extra_cuda_cflags(list[str]|None, optional): Specify additional flags used to compile CUDA files. By default
                                all basic and framework related flags have been included.
                                See `Cuda Compiler Driver NVCC <https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html>`_
                                for details. Default is None.
-        extra_ldflags(list[str], optional): Specify additional flags used to link shared library. See
+        extra_ldflags(list[str]|None, optional): Specify additional flags used to link shared library. See
                                 `GCC Link Options <https://gcc.gnu.org/onlinedocs/gcc/Link-Options.html>`_ for details.
                                 Default is None.
-        extra_include_paths(list[str], optional): Specify additional include path used to search header files. By default
+        extra_include_paths(list[str]|None, optional): Specify additional include path used to search header files. By default
                                 all basic headers are included implicitly from ``site-package/paddle/include`` .
                                 Default is None.
-        extra_library_paths(list[str], optional): Specify additional library path used to search library files. By default
+        extra_library_paths(list[str]|None, optional): Specify additional library path used to search library files. By default
                                 all basic libraries are included implicitly from ``site-packages/paddle/libs`` .
                                 Default is None.
-        build_directory(str, optional): Specify root directory path to put shared library file. If set None,
+        build_directory(str|None, optional): Specify root directory path to put shared library file. If set None,
                             it will use ``PADDLE_EXTENSION_DIR`` from os.environ. Use
                             ``paddle.utils.cpp_extension.get_build_directory()`` to see the location. Default is None.
         verbose(bool, optional): whether to verbose compiled log information. Default is False.
