paddlepaddle--paddle
wehub-resource-sync
643 行
20 KiB
Python
643 行
20 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import importlib
|
||
import inspect
|
||
import logging
|
||
import re
|
||
import sys
|
||
import traceback
|
||
from dataclasses import dataclass
|
||
from functools import cached_property, lru_cache
|
||
from typing import TYPE_CHECKING, Any, Literal, TypeAlias
|
||
|
||
from typing_extensions import get_overloads
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Callable
|
||
from types import ModuleType
|
||
|
||
logging.basicConfig(style="{", format="{message}", level=logging.INFO)
|
||
logger = logging.getLogger("Generating stub file for paddle.Tensor")
|
||
logger.setLevel(logging.INFO)
|
||
|
||
INDENT_SIZE = 4
|
||
INDENT = " " * INDENT_SIZE
|
||
MEANLESS_INDEX = -1
|
||
|
||
|
||
MemberType: TypeAlias = Literal[
|
||
"doc",
|
||
"attribute",
|
||
"method",
|
||
]
|
||
|
||
|
||
@dataclass
|
||
class Member:
|
||
id: int
|
||
name: str
|
||
type: MemberType
|
||
aliases: list[str]
|
||
decorators: list[str]
|
||
signature: str
|
||
doc: str | None
|
||
|
||
def add_alias(self, alias: str):
|
||
self.aliases.append(alias)
|
||
|
||
|
||
@lru_cache
|
||
def _slot_pattern(slot_name: str) -> re.Pattern:
|
||
return re.compile(
|
||
r"(?P<indent> *)#\s*annotation:\s*\$\{" + slot_name + r"\}"
|
||
)
|
||
|
||
|
||
@lru_cache
|
||
def create_builtin_annotation_renamer():
|
||
# NOTE(ooooo-create): Rename built-in types to avoid naming conflicts
|
||
builtin_types = ["int", "bool", "str", "float", "complex", "bytes"]
|
||
regex_string = "|".join([rf"\b{t}\b" for t in builtin_types])
|
||
regex = re.compile(regex_string)
|
||
|
||
def renamer(annotations):
|
||
if annotations is inspect.Signature.empty:
|
||
return annotations
|
||
return regex.sub(lambda m: f"_{m.group(0)}", annotations)
|
||
|
||
return renamer
|
||
|
||
|
||
def rename_builtin_annotation(annotation):
|
||
renamer = create_builtin_annotation_renamer()
|
||
return renamer(annotation)
|
||
|
||
|
||
class TensorGen:
|
||
def __init__(self, template: str = '', prefix: str = 'tensor'):
|
||
self._template = template
|
||
self._template_codes: list[tuple[int, int, str]] = []
|
||
self._prefix = prefix
|
||
|
||
def find_annotation_slot(self, slot_name: str) -> tuple[str, int, int]:
|
||
pattern = _slot_pattern(slot_name)
|
||
slot = []
|
||
for mo in pattern.finditer(self._template):
|
||
_indent = mo.group('indent')
|
||
_start_index, _end_index = mo.span()
|
||
slot.append((_indent, _start_index, _end_index))
|
||
|
||
assert len(slot) == 1, self._template
|
||
return slot[0]
|
||
|
||
@property
|
||
def tensor_docstring(self):
|
||
return self.find_annotation_slot(f'{self._prefix}_docstring')
|
||
|
||
@property
|
||
def tensor_attributes(self):
|
||
return self.find_annotation_slot(f'{self._prefix}_attributes')
|
||
|
||
@property
|
||
def tensor_methods(self):
|
||
return self.find_annotation_slot(f'{self._prefix}_methods')
|
||
|
||
@property
|
||
def tensor_alias(self):
|
||
return self.find_annotation_slot(f'{self._prefix}_alias')
|
||
|
||
@property
|
||
def template_codes(self) -> list[tuple[int, int, str]]:
|
||
return self._template_codes
|
||
|
||
@property
|
||
def template_begin(self) -> int:
|
||
return self._template.find(f'{self._prefix}_begin')
|
||
|
||
@property
|
||
def template_end(self) -> int:
|
||
return self._template.find(f'{self._prefix}_end')
|
||
|
||
def find_apis(self, api_name: str) -> list[dict[str, tuple[str, int, int]]]:
|
||
if self.template_begin < 0 or self.template_end < 0:
|
||
return []
|
||
|
||
pattern = re.compile(
|
||
r"(?P<indent> *)(?P<def_api>def "
|
||
+ api_name
|
||
+ r")(?P<signature>\(.*?\).*?:)(?P<docstring>.*?)(?P<ellipsis>\.{3})(?P<comment>[^\n]*#[^\n]*\n)?",
|
||
re.DOTALL,
|
||
)
|
||
api = []
|
||
for mo in pattern.finditer(self._template):
|
||
_start_index, _end_index = mo.span()
|
||
if not (
|
||
self.template_begin < _start_index < self.template_end
|
||
and self.template_begin < _end_index < self.template_end
|
||
):
|
||
continue
|
||
|
||
_indent = mo.group('indent')
|
||
_signature = mo.group('signature')
|
||
_docstring = mo.group('docstring')
|
||
_ellipsis = mo.group('ellipsis')
|
||
_comment = mo.group('comment')
|
||
_comment = '' if _comment is None else _comment
|
||
|
||
_start_indent, _ = mo.span('indent')
|
||
_start_docstring, _ = mo.span('docstring')
|
||
_, _end_ellipsis = mo.span('ellipsis')
|
||
_start_comment = _end_ellipsis
|
||
_end_comment = _start_comment + len(_comment)
|
||
|
||
assert _start_index == _start_indent
|
||
assert _end_comment == _end_index
|
||
|
||
_api = {
|
||
'indent': (_indent, MEANLESS_INDEX, MEANLESS_INDEX),
|
||
'signature': (_signature, MEANLESS_INDEX, MEANLESS_INDEX),
|
||
'docstring': (_docstring, _start_docstring, _end_comment),
|
||
'ellipsis': (_ellipsis, MEANLESS_INDEX, MEANLESS_INDEX),
|
||
'comment': (_comment, MEANLESS_INDEX, MEANLESS_INDEX),
|
||
}
|
||
api.append(_api)
|
||
|
||
return api
|
||
|
||
def insert_template(self, code: str, start: int, end: int) -> None:
|
||
if start != MEANLESS_INDEX and end != MEANLESS_INDEX:
|
||
self._template_codes.append((start, end, code))
|
||
|
||
def add_method(self, func: Member):
|
||
"""
|
||
1. insert docstring: tensor.prototype.pyi define the method without docstring
|
||
2. insert method: tensor.prototype.pyi NOT define the method
|
||
"""
|
||
methods = self.find_apis(func.name)
|
||
if methods:
|
||
# only use the last method
|
||
method = methods[-1]
|
||
# insert docstring if necessary
|
||
if not method['docstring'][0].strip():
|
||
doc = func.doc
|
||
if doc:
|
||
comment = method['comment'][0]
|
||
doc_start = method['docstring'][1]
|
||
doc_end = method['docstring'][2]
|
||
|
||
api_indent = method['indent'][0]
|
||
|
||
assert len(api_indent) % INDENT_SIZE == 0
|
||
|
||
_indent = api_indent + INDENT
|
||
|
||
_doc = '\n' # new line
|
||
_doc += f'{_indent}r"""\n'
|
||
_doc += with_indent(doc, len(_indent) // INDENT_SIZE)
|
||
_doc += "\n"
|
||
_doc += f'{_indent}"""\n'
|
||
_doc += f'{_indent}...\n'
|
||
_doc += f'{_indent}\n'
|
||
|
||
self.insert_template(comment + _doc, doc_start, doc_end)
|
||
else:
|
||
method_code = '\n'
|
||
for decorator in func.decorators:
|
||
method_code += f"@{decorator}\n"
|
||
|
||
method_code += f"def {func.signature}:\n"
|
||
# 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"
|
||
method_code += f'{INDENT}"""\n'
|
||
method_code += f"{INDENT}...\n"
|
||
|
||
_indent, _, _end_index = self.tensor_methods
|
||
method_code = with_indent(method_code, len(_indent) // INDENT_SIZE)
|
||
self.insert_template(method_code, _end_index, _end_index)
|
||
|
||
def add_alias(self, alias: str, target: str):
|
||
_indent, _, _end_index = self.tensor_alias
|
||
aliases_code = "\n"
|
||
aliases_code += f"{_indent}{alias} = {target}"
|
||
self.insert_template(aliases_code, _end_index, _end_index)
|
||
|
||
def add_attribute(self, name: str, type_: str):
|
||
_indent, _, _end_index = self.tensor_attributes
|
||
attr_code = "\n"
|
||
attr_code += f"{_indent}{name}: {type_}"
|
||
self.insert_template(attr_code, _end_index, _end_index)
|
||
|
||
def add_doc(self, doc: str):
|
||
_indent, _, _end_index = self.tensor_docstring
|
||
docstring = "\n"
|
||
docstring += 'r"""\n'
|
||
docstring += doc
|
||
docstring += "\n"
|
||
docstring += '"""\n'
|
||
docstring = with_indent(docstring, len(_indent) // INDENT_SIZE)
|
||
self.insert_template(docstring, _end_index, _end_index)
|
||
|
||
@classmethod
|
||
def codegen(
|
||
cls, template: str, template_codes: list[tuple[int, int, str]]
|
||
) -> str:
|
||
header = (
|
||
'# This file is auto generated by `tools/gen_tensor_stub.py`.\n\n'
|
||
)
|
||
|
||
_template = []
|
||
start = 0
|
||
end = 0
|
||
for _start, _end, code in sorted(template_codes):
|
||
end = _start
|
||
_template.extend(
|
||
[
|
||
template[start:end],
|
||
code,
|
||
]
|
||
)
|
||
start = _end
|
||
_template.append(template[start:])
|
||
|
||
_content = header + ''.join(_template)
|
||
|
||
return _content
|
||
|
||
|
||
def is_inherited_member(name: str, cls: type) -> bool:
|
||
"""Check if the member is inherited from parent class"""
|
||
|
||
# keep magic methods
|
||
if name.startswith("__") and name.endswith("__"):
|
||
return False
|
||
|
||
if name in cls.__dict__:
|
||
return False
|
||
|
||
for base in cls.__bases__:
|
||
if name in base.__dict__:
|
||
return True
|
||
|
||
return any(is_inherited_member(name, base) for base in cls.__bases__)
|
||
|
||
|
||
def is_property(member: Any) -> bool:
|
||
"""Check if the member is a property"""
|
||
|
||
return isinstance(member, (property, cached_property))
|
||
|
||
|
||
def is_staticmethod(member: Any) -> bool:
|
||
"""Check if the member is a staticmethod"""
|
||
|
||
return isinstance(member, staticmethod)
|
||
|
||
|
||
def is_classmethod(member: Any) -> bool:
|
||
"""Check if the member is a classmethod"""
|
||
|
||
return isinstance(member, classmethod)
|
||
|
||
|
||
def process_lines(code: str, callback: Callable[[str], str]) -> str:
|
||
lines = code.splitlines()
|
||
end_with_newline = code.endswith("\n")
|
||
processed_lines: list[str] = []
|
||
for line in lines:
|
||
processed_lines.append(callback(line))
|
||
processed_code = "\n".join(processed_lines)
|
||
if end_with_newline:
|
||
processed_code += "\n"
|
||
return processed_code
|
||
|
||
|
||
def with_indent(code: str, level: int) -> str:
|
||
def add_indent_line(line: str) -> str:
|
||
if not line:
|
||
return line
|
||
return INDENT + line
|
||
|
||
def remove_indent_line(line: str) -> str:
|
||
if not line:
|
||
return line
|
||
elif line.startswith(INDENT):
|
||
return line[len(INDENT) :]
|
||
else:
|
||
return line
|
||
|
||
if level == 0:
|
||
return code
|
||
elif level > 0:
|
||
if level == 1:
|
||
return process_lines(code, add_indent_line)
|
||
code = process_lines(code, add_indent_line)
|
||
return with_indent(code, level - 1)
|
||
else:
|
||
if level == -1:
|
||
return process_lines(code, remove_indent_line)
|
||
return with_indent(code, level - 1)
|
||
|
||
|
||
def func_sig_to_method_sig(func_sig: str) -> str:
|
||
regex_func_sig = re.compile(
|
||
r"^(?P<method_name>[_a-zA-Z0-9]+)\((?P<arg0>[^,*)]+(:.+)?)?(?P<rest_args>.*)?\)",
|
||
re.DOTALL,
|
||
)
|
||
matched = regex_func_sig.search(func_sig)
|
||
if matched is None:
|
||
# TODO: resolve this case
|
||
logging.warning(f"Cannot parse function signature: {func_sig}")
|
||
return "_(self)"
|
||
|
||
if matched.group('rest_args').startswith('*'):
|
||
method_sig = regex_func_sig.sub(
|
||
r"\g<method_name>(self, \g<rest_args>)", func_sig
|
||
)
|
||
|
||
else:
|
||
method_sig = regex_func_sig.sub(
|
||
r"\g<method_name>(self\g<rest_args>)", func_sig
|
||
)
|
||
|
||
return method_sig
|
||
|
||
|
||
def func_doc_to_method_doc(func_doc: str) -> str:
|
||
# Iterate every line, insert the indent and remove document of the first argument
|
||
method_doc = ""
|
||
is_first_arg = False
|
||
first_arg_offset = 0
|
||
|
||
for line in func_doc.splitlines():
|
||
current_line_offset = len(line) - len(line.lstrip())
|
||
# Remove the first argument (self in Tensor method) from docstring
|
||
if is_first_arg:
|
||
if current_line_offset <= first_arg_offset:
|
||
is_first_arg = False
|
||
if not first_arg_offset:
|
||
first_arg_offset = current_line_offset
|
||
if is_first_arg:
|
||
continue
|
||
method_doc += f"{line}\n" if line else "\n"
|
||
if line.lstrip().startswith("Args:"):
|
||
is_first_arg = True
|
||
|
||
return method_doc
|
||
|
||
|
||
def try_import_paddle() -> ModuleType | None:
|
||
try:
|
||
return importlib.import_module('paddle')
|
||
except ModuleNotFoundError:
|
||
traceback.print_exc(file=sys.stderr)
|
||
sys.stderr.write(
|
||
'''
|
||
ERROR: Can NOT import paddle from `tools/gen_tensor_stub.py` before installation.
|
||
So the stub file `python/paddle/tensor/tensor.pyi` of `paddle.Tensor` may be lost.
|
||
We COULD import paddle without installation with all libs (.dll or .so) copied into dir `paddle/libs`,
|
||
or path already been set for the system. Try the following steps to locate the problem.
|
||
|
||
1. Build with `SKIP_STUB_GEN=ON make -j$(nproc)`.
|
||
2. Install the wheel from `build/python/dist`.
|
||
3. Try to `import paddle` and check the problems.
|
||
'''
|
||
)
|
||
return None
|
||
|
||
|
||
def get_tensor_members(module: str = 'paddle.Tensor') -> dict[int, Member]:
|
||
paddle = try_import_paddle()
|
||
if not paddle:
|
||
raise (
|
||
ModuleNotFoundError(
|
||
'Can NOT import paddle from tools/gen_tensor_stub.py.'
|
||
)
|
||
)
|
||
|
||
tensor_class = eval(module)
|
||
|
||
members: dict[int, Member] = {}
|
||
for name, member in inspect.getmembers(tensor_class):
|
||
member_id = id(member)
|
||
member_doc = inspect.getdoc(member)
|
||
member_doc_cleaned = (
|
||
func_doc_to_method_doc(inspect.cleandoc(member_doc))
|
||
if member_doc is not None
|
||
else None
|
||
)
|
||
try:
|
||
sig = inspect.signature(member)
|
||
sig = sig.replace(
|
||
parameters=[
|
||
p.replace(
|
||
annotation=rename_builtin_annotation(p.annotation)
|
||
)
|
||
for p in sig.parameters.values()
|
||
],
|
||
return_annotation=rename_builtin_annotation(
|
||
sig.return_annotation
|
||
),
|
||
)
|
||
# TODO: classmethod
|
||
member_signature = f"{name}{sig}"
|
||
|
||
except (TypeError, ValueError):
|
||
member_signature = f"{name}()"
|
||
|
||
if is_inherited_member(name, tensor_class):
|
||
continue
|
||
|
||
# Filter out private members except magic methods
|
||
if name.startswith("_") and not (
|
||
name.startswith("__") and name.endswith("__")
|
||
):
|
||
continue
|
||
|
||
if member_id in members:
|
||
members[member_id].add_alias(name)
|
||
continue
|
||
|
||
if name == '__doc__':
|
||
members[member_id] = Member(
|
||
member_id,
|
||
name,
|
||
"doc",
|
||
[],
|
||
[],
|
||
"__doc__",
|
||
member,
|
||
)
|
||
elif is_property(member) or inspect.isdatadescriptor(member):
|
||
members[member_id] = Member(
|
||
member_id,
|
||
name,
|
||
"method",
|
||
[],
|
||
["property"],
|
||
f"{name}(self)",
|
||
member_doc_cleaned,
|
||
)
|
||
elif is_classmethod(member):
|
||
members[member_id] = Member(
|
||
member_id,
|
||
name,
|
||
"method",
|
||
[],
|
||
["classmethod"],
|
||
member_signature,
|
||
member_doc_cleaned,
|
||
)
|
||
elif is_staticmethod(member):
|
||
members[member_id] = Member(
|
||
member_id,
|
||
name,
|
||
"method",
|
||
[],
|
||
["staticmethod"],
|
||
member_signature,
|
||
member_doc_cleaned,
|
||
)
|
||
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)
|
||
_sig = _sig.replace(
|
||
parameters=[
|
||
p.replace(
|
||
annotation=rename_builtin_annotation(p.annotation)
|
||
)
|
||
for p in _sig.parameters.values()
|
||
],
|
||
return_annotation=rename_builtin_annotation(
|
||
_sig.return_annotation
|
||
),
|
||
)
|
||
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,
|
||
"method",
|
||
[],
|
||
[],
|
||
func_sig_to_method_sig(member_signature),
|
||
member_doc_cleaned,
|
||
)
|
||
else:
|
||
logging.debug(f"Skip unknown type of member: {name}, {member}")
|
||
|
||
return members
|
||
|
||
|
||
def get_tensor_template(path: str) -> str:
|
||
with open(path) as f:
|
||
return ''.join(f.readlines())
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser()
|
||
|
||
parser.add_argument(
|
||
"-i",
|
||
"--input-file",
|
||
type=str,
|
||
default="python/paddle/tensor/tensor.prototype.pyi",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--output-file",
|
||
type=str,
|
||
default="python/paddle/tensor/tensor.pyi",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
return args
|
||
|
||
|
||
def generate_stub_file(input_file=None, output_file=None):
|
||
# Get tensor template
|
||
tensor_template = get_tensor_template(input_file)
|
||
|
||
all_members: set[int] = set()
|
||
template_codes: list[tuple[int, int, str]] = []
|
||
# Get members of Tensor
|
||
for module, prefix in [
|
||
['paddle.Tensor', 'tensor'],
|
||
['paddle.base.framework.EagerParamBase', 'eager_param_base'],
|
||
]:
|
||
tensor_members = get_tensor_members(module)
|
||
logging.debug(f'total members in {module}: {len(tensor_members)}')
|
||
|
||
# Generate the Tensor stub
|
||
tensor_gen = TensorGen(tensor_template, prefix)
|
||
for member_id, member in tensor_members.items():
|
||
if member_id in all_members:
|
||
continue
|
||
|
||
if member.type == "method":
|
||
tensor_gen.add_method(member)
|
||
for alias in member.aliases:
|
||
tensor_gen.add_alias(alias, member.name)
|
||
elif member.type == "attribute":
|
||
tensor_gen.add_attribute(member.name, "Any")
|
||
elif member.type == "doc":
|
||
tensor_gen.add_doc(member.doc)
|
||
|
||
all_members |= tensor_members.keys()
|
||
template_codes += tensor_gen.template_codes
|
||
|
||
# Write to target file
|
||
with open(output_file, "w", encoding="utf-8") as f:
|
||
f.write(TensorGen.codegen(tensor_template, template_codes))
|
||
|
||
|
||
def main():
|
||
args = parse_args()
|
||
generate_stub_file(args.input_file, args.output_file)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|