# 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 functools import importlib import inspect import keyword import logging import os import py_compile import re import shutil import tempfile import traceback from pathlib import Path from typing import TYPE_CHECKING, Any, TypedDict import yaml from pybind11_stubgen import ( CLIArgs, Printer, Writer, run, stub_parser_from_args, to_output_and_subdir, ) from pybind11_stubgen.structs import ( Annotation, InvalidExpression, Value, ) if TYPE_CHECKING: from collections.abc import Generator # some invalid attr can NOT be parsed. # to avoid syntax error, we can only do plain replacement. # e.g. {'a': 'b'}, do replace 'a' -> 'b' . BAD_ATTR = {} # add some import modules # e.g. {'a': 'b'}, if not found ' a.' in stub file, # we insert 'b' after 'from __future__ import annotations'. # 'a' should be converted to ' a.' EXTRA_IMPORTS = { 'paddle': 'import paddle', 'typing': 'import typing', 'typing_extensions': 'import typing_extensions', 'pybind11_stubgen': 'import pybind11_stubgen', 'npt': 'import numpy.typing as npt', 'types': 'import types', } # some invalid attr from pybind11. # ref: # - https://pybind11.readthedocs.io/en/latest/advanced/misc.html#avoiding-cpp-types-in-docstrings # - https://pybind11.readthedocs.io/en/latest/advanced/functions.html#default-arguments-revisited # we can add some mappings for conversion, e.g. {'paddle::Tensor': 'paddle.Tensor'} PYBIND11_ATTR_MAPPING = {} # some bad full expression pybind11-stubgen can not catch as invalid exp PYBIND11_INVALID_FULL_MAPPING = { 'PyCodeObject': 'types.CodeType', 'PyInterpreterFrameProxy': 'typing.Any', '_object': 'typing.Any', 'float16': 'numpy.float16', 'Variant': 'typing.Any', 'capsule': 'typing_extensions.CapsuleType', 'cudaDeviceProp': 'typing.Any', 'Tensor': 'paddle.Tensor', 'TensorLike': 'paddle._typing.TensorLike', 'DTypeLike': 'paddle._typing.DTypeLike', 'ShapeLike': 'paddle._typing.ShapeLike', 'Numeric': 'paddle._typing.Numeric', 'TypeGuard': 'typing_extensions.TypeGuard', '_Interpolation': 'paddle.tensor.stat._Interpolation', 'ParamAttrLike': 'paddle._typing.ParamAttrLike', '_POrder': 'paddle.tensor.linalg._POrder', 'TensorOrTensors': 'paddle._typing.TensorOrTensors', } # some bad partial expression pybind11-stubgen can not catch as invalid exp _PYBIND11_INVALID_PART_MAPPING = { 'NestedSequence': 'paddle._typing.NestedSequence', 'Dep': 'Node.Dep', 'Tensor': 'paddle.Tensor', 'TypeGuard': 'typing.TypeGuard', } PYBIND11_INVALID_PART_MAPPING: dict[re.Pattern, str] = { re.compile(rf'(?.*)\(\*args, \*\*kwargs\):') FUNCTION_VALUE_TRANS = { 'true': 'True', 'false': 'False', } # TODO: Duplicate of python/paddle/tensor/tensor.prototype.pyi # Consider a better way to manage these common mappings. OPS_YAML_IMPORTS = """ # Import common typings for generated methods # isort: off from typing import * # noqa: F403 from typing_extensions import * # type: ignore # noqa: F403 from paddle._typing import * # noqa: F403 # isort: on from builtins import ( # noqa: F401 bool as _bool, bytes as _bytes, complex as _complex, float as _float, int as _int, str as _str, ) from collections.abc import Iterator from typing import Any, Literal, overload import numpy.typing as npt import paddle from paddle import ( ParamAttr, # noqa: F401 _typing, ) from paddle.base.dygraph.tensor_patch_methods import ( TensorHookRemoveHelper, # noqa: F401 ) from paddle.tensor.linalg import _POrder # noqa: F401 from paddle.tensor.stat import _Interpolation # noqa: F401 # Special types already defined in tensor.prototype.pyi from paddle import Tensor """ def _get_pybind11_stubgen_annotation_text(annotation: Annotation) -> str: if isinstance(annotation, InvalidExpression): return annotation.text return str(annotation) def _patch_pybind11_invalid_name(): # patch name with suffix '_' if `name` is a keyword like `in` to `in_` def wrap_name(func): @functools.wraps(func) def wrapper(self, arg: Any): if hasattr(arg, 'name') and keyword.iskeyword(arg.name): arg.name += '_' return func(self, arg) return wrapper Printer.print_argument = wrap_name(Printer.print_argument) Printer.print_function = wrap_name(Printer.print_function) def _patch_pybind11_invalid_annotation(): # patch invalid annotation as `Value`, e.g. 'capsule' to 'typing_extensions.CapsuleType' def wrap_name(func): @functools.wraps(func) def wrapper(self, arg: Annotation): _arg_str = _get_pybind11_stubgen_annotation_text(arg) if _arg_str in PYBIND11_INVALID_FULL_MAPPING: arg = Value(PYBIND11_INVALID_FULL_MAPPING[_arg_str]) else: _arg = _arg_str for ( sig_pattern, mapping, ) in PYBIND11_INVALID_PART_MAPPING.items(): _arg = sig_pattern.sub(mapping, _arg) if _arg != _arg_str: arg = Value(_arg) return func(self, arg) return wrapper Printer.print_annotation = wrap_name(Printer.print_annotation) def _patch_pybind11_invalid_exp(): # patch invalid exp with `"xxx"` as a `typing.Any` def print_invalid_exp(self, invalid_expr: InvalidExpression) -> str: _text = invalid_expr.text _text = PYBIND11_ATTR_MAPPING.get(_text, f'"{_text}"') if ( self.invalid_expr_as_ellipses and _text not in PYBIND11_INVALID_FULL_MAPPING.values() ): return "typing.Any" return _text Printer.print_invalid_exp = print_invalid_exp def patch_pybind11_stubgen_printer(): _patch_pybind11_invalid_name() _patch_pybind11_invalid_annotation() _patch_pybind11_invalid_exp() def gen_stub( output_dir: str, module_name: str, ignore_all_errors: bool = False, print_invalid_expressions_as_is: bool = False, ) -> None: logging.basicConfig( level=logging.INFO, format="%(name)s - [%(levelname)7s] %(message)s", ) args = CLIArgs( output_dir=output_dir, root_suffix=None, ignore_invalid_expressions=None, ignore_invalid_identifiers=None, ignore_unresolved_names=None, ignore_all_errors=ignore_all_errors, enum_class_locations=[], numpy_array_wrap_with_annotated=False, numpy_array_use_type_var=False, numpy_array_remove_parameters=False, print_invalid_expressions_as_is=print_invalid_expressions_as_is, print_safe_value_reprs=None, exit_code=False, dry_run=False, stub_extension='pyi', module_name=module_name, ) parser = stub_parser_from_args(args) printer = Printer( invalid_expr_as_ellipses=not args.print_invalid_expressions_as_is ) out_dir, sub_dir = to_output_and_subdir( output_dir=args.output_dir, module_name=args.module_name, root_suffix=args.root_suffix, ) run( parser, printer, args.module_name, out_dir, sub_dir=sub_dir, dry_run=args.dry_run, writer=Writer(stub_ext=args.stub_extension), ) def replace_bad_attr(filename: str): def wrap_text(text): """wrap text to avoid bad math""" return ' ' + text stub_file = None with open(filename, encoding='utf-8') as f: stub_file = f.read() for _bad_attr, _replacement in BAD_ATTR.items(): bad_attr = wrap_text(_bad_attr) replacement = wrap_text(_replacement) stub_file = stub_file.replace(bad_attr, replacement) with open(filename, 'w', encoding='utf-8') as f: f.write(stub_file) def insert_import_modules(filename: str): def wrap_text(text): """wrap text to avoid bad math""" return ' ' + text + '.' future_import = 'from __future__ import annotations\n' stub_file = None with open(filename, encoding='utf-8') as f: stub_file = f.read() for _module, _import_txt in EXTRA_IMPORTS.items(): module = wrap_text(_module) if module in stub_file and _import_txt not in stub_file: stub_file = stub_file.replace( future_import, future_import + _import_txt + '\n' ) with open(filename, 'w', encoding='utf-8') as f: f.write(stub_file) def check_remove_syntax_error(filename: str, limit: int = 10000): """ Args: filename: xxx.pyi limit: the max times try to check syntax error """ pattern_check = re.compile( rf"File.*{re.escape(filename)}.*line (?P\d+)" ) while limit > 0: limit -= 1 # check syntax error err = "" try: py_compile.compile(filename, doraise=True) break except py_compile.PyCompileError as e: err = traceback.format_exc() print(f">>> Syntax error: find syntax error in file: {filename}") # find the line with syntax error match_obj = pattern_check.search(err) if match_obj is not None: line_no = int(match_obj.group("lineno")) # read file source_lines = [] with open(filename, "r", encoding="utf-8") as f: source_lines = f.readlines() del source_lines[line_no - 1] # write new lines with open(filename, "w", encoding="utf-8") as f: f.writelines(source_lines) print( f">>> Syntax error: remove the error line {line_no}, and continue to check ..." ) else: print(">>> Syntax error: no match obj, just continue ...") break def post_process(output_dir: str): """ 1. replace some bad attr 2. check and remove syntax error lines 3. insert some import modules. This should be the last process. """ for root, _, files in os.walk(output_dir): if not files: continue for f in files: # only patch stub files: *.pyi if not f.endswith('.pyi'): continue filename = str(Path(root) / f) # insert modules insert_import_modules(filename) replace_bad_attr(filename) check_remove_syntax_error(filename) # insert modules if necessary insert_import_modules(filename) def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "-o", "--output-dir", type=str, default="python/paddle/", ) parser.add_argument( "-m", "--module-name", type=str, default="", ) parser.add_argument( "--ignore-all-errors", default=False, action="store_true", help="Ignore all errors during module parsing", ) parser.add_argument( "--print-invalid-expressions-as-is", default=False, action="store_true", help="Suppress the replacement with 'typing.Any' of invalid expressions" "found in annotations", ) parser.add_argument( "--ops-yaml", nargs='*', help="Parse ops yaml, the input should be `;` or `;;`" "like `/foo/bar/ops.yaml;paddle.x.y.ops` or /foo/bar/ops.yaml;paddle.x.y.ops;sparse", ) parser.add_argument( "--python-api-info-yaml-path", type=str, default=None, help="the yaml file path for python api info", ) args = parser.parse_args() return args def generate_stub_file( output_dir: str, module_name: str, ignore_all_errors: bool = False, print_invalid_expressions_as_is: bool = False, ops_yaml: list[str] | None = None, python_api_info_yaml_path: str | None = None, ): # patch `pybind11-stubgen` patch_pybind11_stubgen_printer() # generate stub files with tempfile.TemporaryDirectory() as tmpdirname: # gen stub gen_stub( output_dir=tmpdirname, # e.g.: 'Paddle/python/', module_name=module_name, # e.g.: 'paddle.base.libpaddle', ignore_all_errors=ignore_all_errors, print_invalid_expressions_as_is=print_invalid_expressions_as_is, ) # parse ops yaml into file if ops_yaml is not None: ops_yaml_helper = OpsYamlBaseAPI() python_api_info: dict[str, list[str]] = {} if python_api_info_yaml_path is not None: python_api_info = ops_yaml_helper.parse_python_api_info( python_api_info_yaml_path ) for ( yaml_path, dst_module, op_prefix, ) in ops_yaml_helper._parse_yaml_inputs(ops_yaml): dst_module_path = ops_yaml_helper._parse_dst_module_to_path( tmpdirname, module_name, dst_module, ) ops_yaml_helper.parse_yaml_ops( yaml_path, dst_module_path, python_api_info, op_prefix, ) ops_yaml_helper.insert_yaml_imports(dst_module_path) # move stub files into output_dir paths = module_name.split('.') source_path = Path(tmpdirname).joinpath(*paths) if source_path.is_dir(): _path_dst = Path(output_dir).joinpath(paths[-1]) if _path_dst.exists(): shutil.rmtree(str(_path_dst)) else: paths[-1] += '.pyi' _path_dst = Path(output_dir).joinpath(paths[-1]) if _path_dst.exists(): os.remove(str(_path_dst)) shutil.move(str(source_path), output_dir) # post process post_process(output_dir) def load_python_api_function_by_name(name: str) -> Any: components = name.split('.') mod = importlib.import_module(components[0]) fn = mod for comp in components[1:]: fn = getattr(fn, comp) return fn class _OpsYamlInputs(TypedDict): names: list[str] input_info: dict[str, str] class _OpsYamlAttr(TypedDict): names: list[str] attr_info: dict[str, tuple[str, str]] class OpsYamlBaseAPI: """ref: paddle/phi/api/generator/api_base.py""" import_inserted: set[str] = set() # ref: paddle/phi/api/generator/api_base.py def parse_input_and_attr( self, api_name: str, args_config: str, optional_vars: list[str] = [] ) -> tuple[_OpsYamlInputs, _OpsYamlAttr]: inputs = {'names': [], 'input_info': {}} attrs = {'names': [], 'attr_info': {}} args_str = args_config.strip() assert args_str.startswith('(') and args_str.endswith(')'), ( f"Args declaration should start with '(' and end with ')', please check the args of {api_name} in yaml." ) args_str = args_str[1:-1] pattern = re.compile(r',(?![^{]*\})') # support int[] a={1,3} args_list = re.split(pattern, args_str.strip()) args_list = [x.strip() for x in args_list] for item in args_list: item = item.strip() type_and_name = item.split(' ') # match the input tensor has_input = False for in_type_symbol, in_type in INPUT_TYPES_MAP.items(): if type_and_name[0] == in_type_symbol: input_name = type_and_name[1].strip() assert len(input_name) > 0, ( f"The input tensor name should not be empty. Please check the args of {api_name} in yaml." ) assert len(attrs['names']) == 0, ( f"The input Tensor should appear before attributes. please check the position of {api_name}:input({input_name}) in yaml" ) if input_name in optional_vars: in_type = OPTIONAL_TYPES_TRANS[in_type_symbol] inputs['names'].append(input_name) inputs['input_info'][input_name] = in_type has_input = True break if has_input: continue # match the attribute for attr_type_symbol, attr_type in ATTR_TYPES_MAP.items(): if type_and_name[0] == attr_type_symbol: attr_name = item[len(attr_type_symbol) :].strip() assert len(attr_name) > 0, ( f"The attribute name should not be empty. Please check the args of {api_name} in yaml." ) default_value = None if '=' in attr_name: attr_infos = attr_name.split('=') attr_name = attr_infos[0].strip() default_value = attr_infos[1].strip() if attr_name in optional_vars: attr_type = OPTIONAL_TYPES_TRANS[attr_type_symbol] attrs['names'].append(attr_name) attrs['attr_info'][attr_name] = (attr_type, default_value) break return inputs, attrs # ref: paddle/phi/api/generator/api_base.py def parse_output( self, api_name: str, output_config: str ) -> tuple[list[str], Any, Any]: def parse_output_item(output_item): result = re.search( r"(?P[a-zA-Z0-9_[\]]+)\s*(?P\([a-zA-Z0-9_@]+\))?\s*(?P\{[^\}]+\})?", output_item, ) assert result is not None, ( f"{api_name} : the output config parse error." ) out_type = result.group('out_type') assert out_type in OUTPUT_TYPE_MAP, ( f"{api_name} : Output type error: the output type only support Tensor and Tensor[], \ but now is {out_type}." ) out_name = ( 'out' if result.group('name') is None else result.group('name')[1:-1] ) out_size_expr = ( None if result.group('expr') is None else result.group('expr')[1:-1] ) return OUTPUT_TYPE_MAP[out_type], out_name, out_size_expr temp_list = output_config.split(',') if len(temp_list) == 1: out_type, out_name, size_expr = parse_output_item(temp_list[0]) return [out_type], [out_name], [size_expr] else: out_type_list = [] out_name_list = [] out_size_expr_list = [] for output_item in temp_list: out_type, out_name, size_expr = parse_output_item(output_item) out_type_list.append(out_type) out_name_list.append(out_name) out_size_expr_list.append(size_expr) return out_type_list, out_name_list, out_size_expr_list def _make_sig_name(self, name: str) -> str: # 'lambda' -> 'lambda_' if keyword.iskeyword(name): name += '_' return name def _make_attr(self, info: tuple[str, str | None]) -> str: info_name, info_value = info if info_value is None: return info_name if info_name.startswith('list') and '{' in info_value: info_value = info_value.replace('{', '[').replace('}', ']') elif info_value in FUNCTION_VALUE_TRANS: info_value = FUNCTION_VALUE_TRANS[info_value] elif info_name == 'float' and info_value.lower().endswith('f'): info_value = info_value[:-1] elif info_name == 'int' and info_value.lower().endswith('l'): info_value = info_value[:-1] else: try: eval(info_value) except: info_value = f'"{info_value}"' return info_name + ' = ' + info_value def _make_sig(self, name: str, sig: tuple[str, str]) -> str: return self._make_sig_name(name) + ': ' + self._make_attr(sig) def make_op_function( self, name: str, inputs: _OpsYamlInputs, attrs: _OpsYamlAttr, output_type_list: list[str], ) -> str: input_info_names = inputs['names'] input_info = inputs['input_info'] attr_info_names = attrs['names'] attr_info = attrs['attr_info'] _sig_info = [ self._make_sig(_name, (input_info[_name], None)) for _name in input_info_names if _name in input_info ] _sig_attr = [ self._make_sig(_name, attr_info[_name]) for _name in attr_info_names if _name in attr_info ] sig_input = ', '.join(_sig_info + _sig_attr) sig_output = ( output_type_list[0] if len(output_type_list) == 1 else f'tuple[{", ".join(output_type_list)}]' ) return f'def {name}({sig_input}) -> {sig_output}:\n' def make_python_api_function( self, name: str, python_api_names: list[str], ) -> str: fn = load_python_api_function_by_name(python_api_names[0]) sig = inspect.signature(fn) return f'def {name}{sig}:\n' def parse_python_api_info(self, yaml_path: str) -> dict[str, list[str]]: # op name -> python api names # e.g. {'add': ['paddle.add', 'paddle.Tensor.add']} api_info: dict[str, list[str]] = {} with open(yaml_path) as f: api_list = yaml.load(f, Loader=yaml.FullLoader) for api_item_yaml in api_list: op_name = api_item_yaml['op'] api_names = [item.strip() for item in api_item_yaml['name']] api_info[op_name] = api_names return api_info # ref: paddle/phi/api/generator/api_base.py def parse_yaml_ops( self, yaml_file: str, dst_module_path: str, python_api_info: dict[str, list[str]], op_prefix: str | None = None, ) -> None: ops_names = {} ops_file = [] # read stub file generated by pybind11-stubgen and get op names with open(dst_module_path) as f: for line_no, line in enumerate(f.readlines()): match_obj = PATTERN_FUNCTION.match(line) if match_obj is not None: ops_names[match_obj.group('name')] = line_no ops_file.append(line) # read yaml with open(yaml_file) as f: api_list = yaml.load(f, Loader=yaml.FullLoader) for api_item_yaml in api_list: optional_vars = [] if 'optional' in api_item_yaml: optional_vars = [ item.strip() for item in api_item_yaml['optional'].split(',') ] # get op_name, and add op_prefix raw_op_name = api_item_yaml['op'] raw_op_name = ( f'{op_prefix}_{raw_op_name}' if op_prefix is not None else raw_op_name ) op_args = api_item_yaml['args'] op_output = api_item_yaml['output'] # generate input and output op_inputs, op_attrs = self.parse_input_and_attr( raw_op_name, op_args, optional_vars ) output_type_list, _, _ = self.parse_output( raw_op_name, op_output ) # generate full signature from op and inplace op for op_name in [raw_op_name, raw_op_name + '_']: if op_name in ops_names: try: if op_name in python_api_info: op_signature = self.make_python_api_function( op_name, python_api_info[op_name] ) else: op_signature = self.make_op_function( raw_op_name, op_inputs, op_attrs, output_type_list, ) # replace the line from stub file with full signature ops_file[ops_names[op_name]] = op_signature except: print( op_name, op_inputs, op_attrs, output_type_list ) raise with open(dst_module_path, 'w') as f: f.writelines(ops_file) def insert_yaml_imports(self, dst_module_path: str) -> None: # insert imports into file only once if dst_module_path in self.import_inserted: return self.import_inserted.add(dst_module_path) ops_file = [] with open(dst_module_path, 'r') as f: ops_file = f.readlines() import_line_no = 0 for line_no, line in enumerate(ops_file): if line.startswith('from __future__ import annotations'): import_line_no = line_no + 1 break # insert imports ops_file = [ *ops_file[:import_line_no], "\n", *OPS_YAML_IMPORTS.strip().splitlines(keepends=True), "\n", *ops_file[import_line_no:], ] with open(dst_module_path, 'w') as f: f.writelines(ops_file) def _parse_yaml_inputs( self, ops_yaml: list[str], ) -> Generator[tuple[str, str, str], None, None]: for ops in ops_yaml: _ops = ops.split(';') if len(_ops) == 2: yaml_path, dst_module = _ops op_prefix = None elif len(_ops) == 3: yaml_path, dst_module, op_prefix = _ops yield yaml_path, dst_module, op_prefix def _parse_dst_module_to_path( self, base_dir: str, module_name: str, dst_module: str ) -> str: assert dst_module.startswith( module_name ) # e.g.: `paddle.base.libpaddle` in `paddle.base.libpaddle.eager.ops` paths = dst_module.split('.') dst_path = Path(base_dir).joinpath(*paths) if dst_path.is_dir(): dst_path = dst_path.joinpath('__init__.pyi') else: dst_path = dst_path.parent.joinpath(paths[-1] + '.pyi') assert dst_path.exists() return str(dst_path) def main(): args = parse_args() generate_stub_file( output_dir=args.output_dir, module_name=args.module_name, ignore_all_errors=args.ignore_all_errors, print_invalid_expressions_as_is=args.print_invalid_expressions_as_is, ops_yaml=args.ops_yaml, python_api_info_yaml_path=args.python_api_info_yaml_path, ) if __name__ == '__main__': main()