[Typing] 生成 tensor.pyi 时添加 overload 的方法 (#68598)

megemini · web-flow · commit 3865ce6ea3a3 · 2024-10-11T00:07:19.000+08:00
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -17,10 +17,11 @@
 import hashlib
 import inspect
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, overload
+from typing import TYPE_CHECKING, Any, Callable
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import overload
 
 import paddle
 from paddle import _C_ops, profiler
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
@@ -13,10 +13,10 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, overload
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, overload
 
 import paddle
 from paddle import _C_ops
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -16,9 +16,10 @@
 
 import functools
 import math
-from typing import TYPE_CHECKING, Any, Literal, overload
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
+from typing_extensions import overload
 
 import paddle
 from paddle import _C_ops
diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -14,9 +14,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, overload
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
+from typing_extensions import overload
 
 import paddle
 from paddle import _C_ops
diff --git a/python/paddle/tensor/stat.py b/python/paddle/tensor/stat.py
@@ -14,9 +14,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, overload
+from typing import TYPE_CHECKING, Literal
 
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, overload
 
 import paddle
 from paddle import _C_ops
diff --git a/tools/gen_tensor_stub.py b/tools/gen_tensor_stub.py
@@ -25,7 +25,7 @@
 from functools import cached_property, lru_cache
 from typing import TYPE_CHECKING, Any, Callable, Literal, Protocol
 
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, get_overloads
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -221,7 +221,9 @@ def add_method(self, func: Member):
                 method_code += f"@{decorator}\n"
 
             method_code += f"def {func.signature}:\n"
-            if func.doc:
+            # do NOT insert docs from overload methods,
+            # because we always add a plain method
+            if func.doc and func.decorators != ["overload"]:
                 method_code += f'{INDENT}r"""\n'
                 method_code += with_indent(func.doc, 1)
                 method_code += "\n"
@@ -506,11 +508,34 @@ def get_tensor_members(module: str = 'paddle.Tensor') -> dict[int, Member]:
                 member_signature,
                 member_doc_cleaned,
             )
-        elif (
-            inspect.isfunction(member)
-            or inspect.ismethod(member)
-            or inspect.ismethoddescriptor(member)
-        ):
+        elif inspect.isfunction(member) or inspect.ismethod(member):
+            # `all_signatures`： list[[member id, decorators, signature]]
+            # with atleast an original method
+            all_signatures = [[member_id, [], member_signature]]
+
+            # try to get overloads
+            _overloads = get_overloads(member)
+            for f in _overloads:
+                _sig = inspect.signature(f)
+                all_signatures.append(
+                    [
+                        id(f),
+                        ["overload"],
+                        f"{name}{_sig}".replace("Ellipsis", "..."),
+                    ]
+                )
+
+            for _member_id, _decorators, _sig in all_signatures:
+                members[_member_id] = Member(
+                    _member_id,
+                    name,
+                    "method",
+                    [],
+                    _decorators,
+                    func_sig_to_method_sig(_sig),
+                    member_doc_cleaned,
+                )
+        elif inspect.ismethoddescriptor(member):
             members[member_id] = Member(
                 member_id,
                 name,
@@ -522,6 +547,7 @@ def get_tensor_members(module: str = 'paddle.Tensor') -> dict[int, Member]:
             )
         else:
             logging.debug(f"Skip unknown type of member: {name}, {member}")
+
     return members
 
 
