gemini-code-assist[bot] commented on code in PR #268:
URL: https://github.com/apache/tvm-ffi/pull/268#discussion_r2532309160
##########
python/tvm_ffi/stub/utils.py:
##########
@@ -30,3 +34,124 @@ class Options:
files: list[str] = dataclasses.field(default_factory=list)
verbose: bool = False
dry_run: bool = False
+
+
[email protected](init=False)
+class NamedTypeSchema(TypeSchema):
+ """A type schema with an associated name."""
+
+ name: str
+
+ def __init__(self, name: str, schema: TypeSchema) -> None:
+ """Initialize a `NamedTypeSchema` with the given name and type
schema."""
+ super().__init__( # type: ignore[call-arg]
+ origin=schema.origin,
+ args=schema.args,
+ )
+ self.name = name
+
+
[email protected]
+class FuncInfo:
+ """Information of a function."""
+
+ schema: NamedTypeSchema
+ is_member: bool
+
+ @staticmethod
+ def from_global_name(name: str) -> FuncInfo:
+ """Construct a `FuncInfo` from a string name of this global
function."""
+ from tvm_ffi.registry import get_global_func_metadata # noqa: PLC0415
+
+ return FuncInfo(
+ schema=NamedTypeSchema(
+ name=name,
+
schema=TypeSchema.from_json_str(get_global_func_metadata(name)["type_schema"]),
+ ),
+ is_member=False,
+ )
+
+ def gen(self, ty_map: Callable[[str], str], indent: int) -> str:
+ """Generate a function signature string for this function."""
+ try:
+ _, func_name = self.schema.name.rsplit(".", 1)
+ except ValueError:
+ func_name = self.schema.name
+ buf = StringIO()
+ buf.write(" " * indent)
+ buf.write(f"def {func_name}(")
+ if self.schema.origin != "Callable":
+ raise ValueError(f"Expected Callable type schema, but got:
{self.schema}")
+ if not self.schema.args:
+ ty_map("Any")
+ buf.write("*args: Any) -> Any: ...")
+ return buf.getvalue()
+ arg_ret = self.schema.args[0]
+ arg_args = self.schema.args[1:]
+ for i, arg in enumerate(arg_args):
+ if self.is_member and i == 0:
+ buf.write("self, ")
+ else:
+ buf.write(f"_{i}: ")
+ buf.write(arg.repr(ty_map))
+ buf.write(", ")
+ if arg_args:
+ buf.write("/")
+ buf.write(") -> ")
+ buf.write(arg_ret.repr(ty_map))
+ buf.write(": ...")
+ return buf.getvalue()
+
+
[email protected]
+class ObjectInfo:
+ """Information of an object type, including its fields and methods."""
+
+ fields: list[NamedTypeSchema]
+ methods: list[FuncInfo]
+
+ @staticmethod
+ def from_type_key(type_key: str) -> ObjectInfo:
+ """Construct an `ObjectInfo` from a type key."""
+ from tvm_ffi.core import _lookup_or_register_type_info_from_type_key
# noqa: PLC0415
+
+ type_info = _lookup_or_register_type_info_from_type_key(type_key)
+ return ObjectInfo(
+ fields=[
+ NamedTypeSchema(
+ name=field.name,
+
schema=TypeSchema.from_json_str(field.metadata["type_schema"]),
+ )
+ for field in type_info.fields
+ ],
+ methods=[
+ FuncInfo(
+ schema=NamedTypeSchema(
+ name=FN_NAME_MAP.get(method.name, method.name),
+
schema=TypeSchema.from_json_str(method.metadata["type_schema"]),
+ ),
+ is_member=not method.is_static,
+ )
+ for method in type_info.methods
+ ],
+ )
+
+ def gen_fields(self, ty_map: Callable[[str], str], indent: int) ->
list[str]:
+ """Generate field definitions for this object."""
+ indent_str = " " * indent
+ return [f"{indent_str}{field.name}: {field.repr(ty_map)}" for field in
self.fields]
+
+ def gen_methods(self, ty_map: Callable[[str], str], indent: int) ->
list[str]:
+ """Generate method definitions for this object."""
+ indent_str = " " * indent
+ ret = []
+ for method in self.methods:
+ if not method.is_member:
+ ret.append(f"{indent_str}@staticmethod")
+ ret.append(method.gen(ty_map, indent))
+ return ret
+
+
+FN_NAME_MAP = {
+ "__ffi_init__": "__c_ffi_init__",
+}
Review Comment:
It's a good practice to define constants at the top of the file, before they
are used. `FN_NAME_MAP` is used in `ObjectInfo.from_type_key` on line 130.
Moving this definition to the top of the file would improve readability and
maintainability.
##########
python/tvm_ffi/stub/utils.py:
##########
@@ -30,3 +34,124 @@ class Options:
files: list[str] = dataclasses.field(default_factory=list)
verbose: bool = False
dry_run: bool = False
+
+
[email protected](init=False)
+class NamedTypeSchema(TypeSchema):
+ """A type schema with an associated name."""
+
+ name: str
+
+ def __init__(self, name: str, schema: TypeSchema) -> None:
+ """Initialize a `NamedTypeSchema` with the given name and type
schema."""
+ super().__init__( # type: ignore[call-arg]
Review Comment:
This `type: ignore[call-arg]` might no longer be necessary. This pull
request adds an `__init__` definition to `TypeSchema` in `core.pyi`, so type
checkers should now be able to correctly validate this `super().__init__` call.
Consider removing the ignore directive.
```suggestion
super().__init__(
```
##########
python/tvm_ffi/stub/codegen.py:
##########
@@ -181,3 +120,19 @@ def _make_line(module: str, names: list[str], indent: int)
-> str:
"# fmt: on",
code.lines[-1],
]
+
+
+def generate_all(code: CodeBlock, names: set[str], opt: Options) -> None:
+ """Generate an `__all__` variable for the given names."""
+ del opt
Review Comment:
While `del opt` works to silence unused argument warnings, a more idiomatic
Python convention is to prefix the unused argument with an underscore (e.g.,
`_opt`). This signals the intent that the argument is intentionally unused.
```suggestion
def generate_all(code: CodeBlock, names: set[str], _opt: Options) -> None:
"""Generate an `__all__` variable for the given names."""
```
##########
python/tvm_ffi/stub/cli.py:
##########
@@ -56,17 +56,37 @@ def __main__() -> int:
"""
opt = _parse_args()
dlls = [ctypes.CDLL(lib) for lib in opt.dlls]
- global_funcs = collect_global_funcs()
files: list[FileInfo] = collect_files([Path(f) for f in opt.files])
# Stage 1: Process `tvm-ffi-stubgen(ty-map)`
- ty_map: dict[str, str] = collect_ty_maps(files, opt)
+ ty_map: dict[str, str] = C.TY_MAP_DEFAULTS.copy()
+
+ def _stage_1(file: FileInfo) -> None:
+ for code in file.code_blocks:
+ if code.kind == "ty-map":
+ try:
+ lhs, rhs = code.param.split("->")
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid ty_map format at line {code.lineno_start}.
Example: `A.B -> C.D`"
+ ) from e
+ ty_map[lhs.strip()] = rhs.strip()
+
+ for file in files:
+ try:
+ _stage_1(file)
+ except:
Review Comment:
Using a bare `except:` is generally discouraged as it catches exceptions you
probably don't intend to, such as `SystemExit` or `KeyboardInterrupt`. It's
better to explicitly catch `Exception` to handle expected errors while letting
system-level exceptions propagate.
```suggestion
except Exception:
```
--
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]
