gemini-code-assist[bot] commented on code in PR #608:
URL: https://github.com/apache/tvm-ffi/pull/608#discussion_r3365500790
##########
python/tvm_ffi/stub/python_generator/codegen.py:
##########
@@ -257,20 +318,20 @@ def generate_ffi_api(
if not code_blocks:
append += f"""\"\"\"FFI API bindings for {module_name}.\"\"\"\n"""
if not any(code.kind == "import-section" for code in code_blocks):
- append += C.PROMPT_IMPORT_SECTION
+ append += _prompt_import_section(syntax)
# Part 1. Library loading
if is_root:
- append += C._prompt_import_object("tvm_ffi.libinfo.load_lib_module",
"_FFI_LOAD_LIB")
+ append += _prompt_import_object("tvm_ffi.libinfo.load_lib_module",
"_FFI_LOAD_LIB", syntax)
append += f"""LIB = _FFI_LOAD_LIB("{init_cfg.pkg}",
"{init_cfg.shared_target}")\n"""
# Part 2. Global functions
if not any(code.kind == "global" for code in code_blocks):
- append += C._prompt_globals(module_name)
+ append += _prompt_globals(module_name, syntax)
# Part 3. Object types
if object_infos:
- append += C._prompt_import_object("tvm_ffi.register_object",
"_FFI_REG_OBJ")
+ append += _prompt_import_object("tvm_ffi.register_object",
"_FFI_REG_OBJ", syntax)
defined_type_keys = {info.type_key for info in object_infos if
info.type_key}
Review Comment:
The `defined_type_keys` set is constructed using the unmapped
`info.type_key` values. However, during the loop, both `type_key` and
`parent_type_key` are mapped using `ty_map`. This discrepancy can cause
`parent_type_key not in defined_type_keys` to incorrectly evaluate to `True`
for parent types that are actually defined in the same module but have been
mapped via `ty_map`. We should construct `defined_type_keys` using the mapped
type keys.
```suggestion
defined_type_keys = {ty_map.get(info.type_key, info.type_key) for info
in object_infos if info.type_key}
```
##########
python/tvm_ffi/stub/python_generator/generator.py:
##########
@@ -0,0 +1,167 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+"""The Python code generator for ``tvm-ffi-stubgen``.
+
+:class:`PythonGenerator` implements the
:class:`tvm_ffi.stub.generator.Generator`
+protocol by delegating to :mod:`.codegen`. It owns the Python notion of an
+import (:class:`.utils.ImportItem` / :class:`.utils.PythonImports`); the
+language-agnostic pipeline only ever sees the opaque collector.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+from .. import consts as C
+from . import codegen as G
+from . import consts as PC
+from .utils import ImportItem, PythonImports
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ from tvm_ffi.core import TypeSchema
+
+ from ..file_utils import CodeBlock
+ from ..utils import FuncInfo, InitConfig, ObjectInfo, Options
+
+
+class PythonGenerator:
+ """Generator that emits Python type stubs by delegating to
:mod:`.codegen`."""
+
+ name = "python"
+ syntax = C.PYTHON_SYNTAX
+
+ def default_ty_map(self) -> dict[str, str]:
+ """Return the default FFI-origin -> Python-type name map."""
+ return PC.TY_MAP_DEFAULTS.copy()
+
+ def render_type(self, schema: TypeSchema, ty_render: Callable[[str], str])
-> str:
+ """Render a type schema using Python typing syntax (delegates to
`TypeSchema.repr`)."""
+ return schema.repr(ty_render)
+
+ # --- import collection (Python representation is private) ---------------
+
+ def new_imports(self) -> PythonImports:
+ """Create an empty import collector."""
+ return PythonImports()
+
+ def add_imported_object(
+ self, imports: PythonImports, name: str, type_checking_only: str,
alias: str
+ ) -> None:
+ """Record an ``import-object`` directive into the collector."""
+ tco = (
+ bool(type_checking_only)
+ and isinstance(type_checking_only, str)
+ and type_checking_only.lower() == "true"
+ )
Review Comment:
The parsing of `type_checking_only` assumes it is always a string, but if a
boolean is passed (e.g., in programmatic usage or future extensions),
`isinstance(type_checking_only, str)` will evaluate to `False`, resulting in
`tco` being incorrectly set to `False`. We can make this check more robust and
simpler by handling both boolean and string types directly.
```python
if isinstance(type_checking_only, bool):
tco = type_checking_only
else:
tco = str(type_checking_only).lower() == "true"
```
##########
python/tvm_ffi/stub/python_generator/utils.py:
##########
@@ -0,0 +1,306 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+"""Python generator helpers for ``tvm-ffi-stubgen``.
+
+This module groups two Python-specific concerns:
+
+- import modelling (:class:`ImportItem`, :class:`PythonImports`)
+- stub rendering helpers for function/object signatures
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from io import StringIO
+from typing import Callable
+
+from ..utils import FuncInfo, ObjectInfo
+from . import consts as C
+
+
[email protected](frozen=True, eq=True)
+class ImportItem:
+ """An import statement item."""
+
+ mod: str
+ name: str
+ type_checking_only: bool = False
+ alias: str | None = None
+
+ def __init__(
+ self,
+ name: str,
+ type_checking_only: bool = False,
+ alias: str | None = None,
+ ) -> None:
+ """Initialize an `ImportItem` with the given module name and optional
alias."""
+ if "." in name:
+ mod, name = name.rsplit(".", 1)
+ for mod_prefix, mod_replacement in C.MOD_MAP.items():
+ if mod.startswith(mod_prefix):
+ mod = mod.replace(mod_prefix, mod_replacement, 1)
+ break
Review Comment:
The prefix matching for module rewrites uses `mod.startswith(mod_prefix)`.
This can lead to incorrect rewrites if a module name starts with the prefix but
is not a sub-module (for example, a module named `ffi_control` would match the
`ffi` prefix and be incorrectly rewritten to `tvm_ffi_control`). We should
ensure that we only match exact module names or sub-modules by checking for a
dot boundary.
```suggestion
for mod_prefix, mod_replacement in C.MOD_MAP.items():
if mod == mod_prefix or mod.startswith(mod_prefix + "."):
mod = mod.replace(mod_prefix, mod_replacement, 1)
break
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
