From 0a17cda6bbd5fb3ba1abd4aaca6ecc4561116abb Mon Sep 17 00:00:00 2001 From: baominghelly <1508269885@qq.com> Date: Fri, 10 Jul 2026 06:36:27 +0000 Subject: [PATCH] Add generate wrapper change --- scripts/generate_wrappers.py | 132 ++++++++++++++++++++++++++++---- tests/test_generate_wrappers.py | 65 ++++++++++++++++ 2 files changed, 181 insertions(+), 16 deletions(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index e2630ad26..f4d0ea1cd 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -242,6 +242,12 @@ def __call__(self, op_name): elif node.kind == CursorKind.CXX_METHOD and node.spelling == "operator()": calls.append(node) + parsed_operator = _parse_operator_header(op_name) + constructors = _prefer_header_type_spellings( + constructors, parsed_operator.constructors + ) + calls = _prefer_header_type_spellings(calls, parsed_operator.calls) + return _Operator(op_name, constructors, calls) @staticmethod @@ -275,6 +281,27 @@ def _parse_operator_header(op_name): return _Operator(op_name, constructors, calls) +def _prefer_header_type_spellings(functions, header_functions): + header_by_names = {} + + for function in header_functions: + names = tuple(arg.spelling for arg in function.get_arguments()) + header_by_names.setdefault(names, []).append(function) + + aligned = [] + + for function in functions: + names = tuple(arg.spelling for arg in function.get_arguments()) + matches = header_by_names.get(names, ()) + + if len(matches) == 1: + aligned.append(matches[0]) + else: + aligned.append(function) + + return aligned + + def _strip_cpp_comments(source): source = re.sub(r"/\*.*?\*/", "", source, flags=re.DOTALL) return re.sub(r"//.*", "", source) @@ -443,6 +470,22 @@ def _find_vector_int64_params(op_name): return set(re.findall(r"std::vector\s+(\w+)", source)) +def _find_optional_vector_int64_params(op_name): + """Return params declared as `std::optional>`. + + libclang can report these as plain `int`, just like non-optional vector + params, so scan the base header to preserve the generated API signature. + """ + source = _find_base_header(op_name).read_text() + + return set( + re.findall( + r"std::optional\s*<\s*std::vector\s*<\s*int64_t\s*>\s*>\s+(\w+)", + source, + ) + ) + + def _find_tensor_params(op_name): source = _find_base_header(op_name).read_text() @@ -474,6 +517,7 @@ def _generate_pybind11(operator): optional_non_tensor_params = _find_optional_non_tensor_params(operator.name) vector_tensor_params = _find_vector_tensor_params(operator.name) vector_int64_params = _find_vector_int64_params(operator.name) + optional_vector_int64_params = _find_optional_vector_int64_params(operator.name) params_with_defaults = _find_params_with_defaults(operator.name) def _is_optional_tensor(arg): @@ -490,6 +534,9 @@ def _is_optional_tensor(arg): def _is_optional(arg): return "std::optional" in arg.type.spelling + def _is_optional_vector_int64(arg): + return arg.spelling in optional_vector_int64_params + def _is_vector_tensor(arg): if arg.spelling in vector_tensor_params: return True @@ -510,6 +557,10 @@ def _generate_params(node): parts.append(f"std::optional {arg.spelling}") elif _is_vector_tensor(arg): parts.append(f"std::vector {arg.spelling}") + elif _is_optional_vector_int64(arg): + parts.append( + f"const std::optional> {arg.spelling}" + ) elif _is_vector_int64(arg): parts.append(f"const std::vector {arg.spelling}") else: @@ -908,6 +959,7 @@ def _generate_generated_dispatch_entries(operator): tensor_params = _find_tensor_params(operator.name) vector_tensor_params = _find_vector_tensor_params(operator.name) vector_int64_params = _find_vector_int64_params(operator.name) + optional_vector_int64_params = _find_optional_vector_int64_params(operator.name) def _is_optional_tensor(arg): spelling = arg.type.spelling @@ -920,6 +972,9 @@ def _is_optional_tensor(arg): return arg.spelling in optional_tensor_params + def _is_optional_vector_int64(arg): + return arg.spelling in optional_vector_int64_params + def _is_vector_tensor(arg): if arg.spelling in vector_tensor_params: return True @@ -933,10 +988,13 @@ def _is_tensor(arg): if arg.spelling in optional_non_tensor_params: return False - if arg.spelling in tensor_params: + if "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling: return True - return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling + if _is_known_non_tensor_type(arg.type.spelling): + return False + + return arg.spelling in tensor_params def _generate_params(node): parts = [] @@ -949,6 +1007,8 @@ def _generate_params(node): parts.append(f"std::optional {arg.spelling}") elif _is_vector_tensor(arg): parts.append(f"std::vector {arg.spelling}") + elif _is_optional_vector_int64(arg): + parts.append(f"std::optional> {arg.spelling}") elif _is_vector_int64(arg): parts.append(f"std::vector {arg.spelling}") elif _is_tensor(arg): @@ -1108,12 +1168,32 @@ def _strip_top_level_const(type_spelling): return type_spelling +def _is_known_non_tensor_type(type_spelling): + type_spelling = _strip_top_level_const(type_spelling) + + if "Tensor" in type_spelling or "TensorView" in type_spelling: + return False + + if type_spelling in { + "bool", + "double", + "float", + "int64_t", + "std::size_t", + "std::string", + }: + return True + + return type_spelling.startswith("std::optional<") + + def _generate_operator_call_instantiation_entries(operator): optional_tensor_params = _find_optional_tensor_params(operator.name) optional_non_tensor_params = _find_optional_non_tensor_params(operator.name) tensor_params = _find_tensor_params(operator.name) vector_tensor_params = _find_vector_tensor_params(operator.name) vector_int64_params = _find_vector_int64_params(operator.name) + optional_vector_int64_params = _find_optional_vector_int64_params(operator.name) def _is_optional_tensor(arg): spelling = arg.type.spelling @@ -1129,6 +1209,9 @@ def _is_optional_tensor(arg): return False + def _is_optional_vector_int64(arg): + return arg.spelling in optional_vector_int64_params + def _is_vector_tensor(arg): if arg.spelling in vector_tensor_params: return True @@ -1144,10 +1227,13 @@ def _is_tensor(arg): if arg.spelling in optional_non_tensor_params: return False - if arg.spelling in tensor_params: + if "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling: return True - return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling + if _is_known_non_tensor_type(arg.type.spelling): + return False + + return arg.spelling in tensor_params def _normalized_type(arg): if _is_optional_tensor(arg): @@ -1156,6 +1242,9 @@ def _normalized_type(arg): if _is_vector_tensor(arg): return "std::vector" + if _is_optional_vector_int64(arg): + return "std::optional>" + if _is_vector_int64(arg): return "std::vector" @@ -1188,6 +1277,16 @@ def _append_optional_params(prefix, params): declarations = [] definitions = [] + seen_entries = set() + + def _append_unique(declaration, definition): + if definition in seen_entries: + return + + seen_entries.add(definition) + declarations.append(declaration) + definitions.append(definition) + for call in operator.calls: template_arguments = _generate_template_arguments(call) params = _generate_parameters(call) @@ -1202,14 +1301,13 @@ def _append_optional_params(prefix, params): ) make_template_arguments = ", ".join( - f"const {_strip_top_level_const(arg.type.spelling)}&" for arg in rest_args + f"const {_normalized_type(arg)}&" for arg in rest_args ) make_params = ", ".join( - f"const {_strip_top_level_const(arg.type.spelling)}& {arg.spelling}" - for arg in rest_args + f"const {_normalized_type(arg)}& {arg.spelling}" for arg in rest_args ) make_function_params = _append_optional_params( - f"const Config& config, const {_strip_top_level_const(first_arg.type.spelling)} {first_arg.spelling}", + f"const Config& config, const {_normalized_type(first_arg)} {first_arg.spelling}", make_params, ) make_instantiation = f"Operator<{op_type}>::Make<{make_template_arguments}>({make_function_params})" @@ -1218,16 +1316,18 @@ def _append_optional_params(prefix, params): f"Operator<{op_type}>::operator()<{template_arguments}>({params}) const" ) - declarations.append(f"extern template void {instantiation};") - declarations.append( - f"extern template std::unique_ptr> {make_instantiation};" + _append_unique( + f"extern template void {instantiation};", + f"template void {instantiation};", ) - declarations.append(f"extern template void {operator_instantiation};") - definitions.append(f"template void {instantiation};") - definitions.append( - f"template std::unique_ptr> {make_instantiation};" + _append_unique( + f"extern template std::unique_ptr> {make_instantiation};", + f"template std::unique_ptr> {make_instantiation};", + ) + _append_unique( + f"extern template void {operator_instantiation};", + f"template void {operator_instantiation};", ) - definitions.append(f"template void {operator_instantiation};") return declarations, definitions diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index e571f6c90..66835e1f2 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -125,6 +125,71 @@ class Clamp { ) in text +def test_extractor_prefers_header_types_for_reused_parameter_names( + monkeypatch, tmp_path +): + module = _load_generator_module() + base_header = tmp_path / "histogram.h" + base_header.write_text( + """ +class Histogram { + public: + virtual void operator()(const Tensor input, const Tensor bins, + const std::optional weight, + const bool density, Tensor hist, + Tensor bin_edges) const = 0; + virtual void operator()(const Tensor input, const int64_t bins, + const std::optional> range, + const std::optional weight, + const bool density, Tensor hist, + Tensor bin_edges) const = 0; +}; +""" + ) + monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header) + + clang_calls = [ + module._ParsedFunction( + [ + module._ParsedArgument("const int", "input"), + module._ParsedArgument("const int", "bins"), + module._ParsedArgument("const int", "weight"), + module._ParsedArgument("const bool", "density"), + module._ParsedArgument("int", "hist"), + module._ParsedArgument("int", "bin_edges"), + ] + ), + module._ParsedFunction( + [ + module._ParsedArgument("const int", "input"), + module._ParsedArgument("const int64_t", "bins"), + module._ParsedArgument("const int", "range"), + module._ParsedArgument("const int", "weight"), + module._ParsedArgument("const bool", "density"), + module._ParsedArgument("int", "hist"), + module._ParsedArgument("int", "bin_edges"), + ] + ), + ] + + operator = module._Operator( + "histogram", + constructors=[], + calls=module._prefer_header_type_spellings( + clang_calls, module._parse_operator_header("histogram").calls + ), + ) + + declarations, _ = module._generate_operator_call_instantiation_entries(operator) + text = "\n".join(declarations) + + assert ( + "Call>, " + "std::optional, bool, Tensor, Tensor>" + ) in text + assert "Call" not in text + + def test_pybind_default_implementation_uses_first_active_index(monkeypatch, tmp_path): module = _load_generator_module() base_header = tmp_path / "mul.h"