gemini-code-assist[bot] commented on code in PR #295:
URL: https://github.com/apache/tvm-ffi/pull/295#discussion_r2583959842


##########
python/tvm_ffi/stub/analysis.py:
##########
@@ -34,8 +35,27 @@ def collect_global_funcs() -> dict[str, list[FuncInfo]]:
         except ValueError:
             print(f"{C.TERM_YELLOW}[Skipped] Invalid name in global function: 
{name}{C.TERM_RESET}")
         else:
-            global_funcs.setdefault(prefix, 
[]).append(FuncInfo.from_global_name(name))
+            try:
+                global_funcs.setdefault(prefix, 
[]).append(FuncInfo.from_global_name(name))
+            except:

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   Using a bare `except:` is generally discouraged as it can catch exceptions 
you don't intend to, like `SystemExit` or `KeyboardInterrupt`. It's better to 
be explicit about the exceptions you want to catch. In this case, `except 
Exception:` would be a safer choice to handle errors during `FuncInfo` creation 
without masking critical system-level exceptions.
   
   ```suggestion
               except Exception:
   ```



##########
python/tvm_ffi/stub/cli.py:
##########
@@ -55,71 +55,43 @@ def __main__() -> int:
     overview and examples of the block syntax.
     """
     opt = _parse_args()
+    for imp in opt.imports or []:
+        importlib.import_module(imp)
+    if opt.init_path:
+        opt.files.append(opt.init_path)
     dlls = [ctypes.CDLL(lib) for lib in opt.dlls]
     files: list[FileInfo] = collect_files([Path(f) for f in opt.files])
+    global_funcs: dict[str, list[FuncInfo]] = collect_global_funcs()
 
-    # Stage 1: Process `tvm-ffi-stubgen(ty-map)`
+    # Stage 1: Collect information
+    # - type maps: `tvm-ffi-stubgen(ty-map)`
+    # - defined global functions: `tvm-ffi-stubgen(begin): global/...`
+    # - defined object types: `tvm-ffi-stubgen(begin): object/...`
     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)
+            _stage_1(file, ty_map)
         except Exception:
             print(
                 f'{C.TERM_RED}[Failed] File "{file.path}": 
{traceback.format_exc()}{C.TERM_RESET}'
             )
 
-    # Stage 2: Process
+    # Stage 2. Generate stubs if they are not defined on the file.
+    if opt.init_path:
+        _stage_2(
+            files,
+            init_path=Path(opt.init_path).resolve(),
+            global_funcs=global_funcs,
+        )
+
+    # Stage 3: Process
     # - `tvm-ffi-stubgen(begin): global/...`
     # - `tvm-ffi-stubgen(begin): object/...`
-    global_funcs = collect_global_funcs()
-
-    def _stage_2(file: FileInfo) -> None:
-        all_defined = set()
+    for file in files:
         if opt.verbose:
             print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}")
-        ty_used: set[str] = set()
-        ty_on_file: set[str] = set()
-        fn_ty_map_fn = _fn_ty_map(ty_map, ty_used)
-        # Stage 2.1. Process `tvm-ffi-stubgen(begin): global/...`
-        for code in file.code_blocks:
-            if code.kind == "global":
-                funcs = global_funcs.get(code.param, [])
-                for func in funcs:
-                    all_defined.add(func.schema.name)
-                G.generate_global_funcs(code, funcs, fn_ty_map_fn, opt)
-        # Stage 2.2. Process `tvm-ffi-stubgen(begin): object/...`
-        for code in file.code_blocks:
-            if code.kind == "object":
-                type_key = code.param
-                ty_on_file.add(ty_map.get(type_key, type_key))
-                G.generate_object(code, fn_ty_map_fn, opt)
-        # Stage 2.3. Add imports for used types.
-        for code in file.code_blocks:
-            if code.kind == "import":
-                G.generate_imports(code, ty_used - ty_on_file, opt)
-                break  # Only one import block per file is supported for now.
-        # Stage 2.4. Add `__all__` for defined classes and functions.
-        for code in file.code_blocks:
-            if code.kind == "__all__":
-                G.generate_all(code, all_defined | ty_on_file, opt)
-                break  # Only one __all__ block per file is supported for now.
-        file.update(show_diff=opt.verbose, dry_run=opt.dry_run)
-
-    for file in files:
         try:
-            _stage_2(file)
+            _stage_3(file, opt, ty_map, global_funcs)
         except:

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   Using a bare `except:` is risky because it catches all exceptions, including 
system-exiting ones like `SystemExit` and `KeyboardInterrupt`. It's safer to 
catch `Exception` to handle application-level errors without interfering with 
process termination signals.
   
   ```suggestion
           except Exception:
   ```



##########
python/tvm_ffi/stub/cli.py:
##########
@@ -128,6 +100,122 @@ def _stage_2(file: FileInfo) -> None:
     return 0
 
 
+def _stage_1(
+    file: FileInfo,
+    ty_map: dict[str, str],
+) -> None:
+    for code in file.code_blocks:
+        if code.kind == "ty-map":
+            try:
+                assert isinstance(code.param, str)
+                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()
+
+
+def _stage_2(
+    files: list[FileInfo],
+    init_path: Path,
+    global_funcs: dict[str, list[FuncInfo]],
+) -> None:
+    def _find_or_insert_file(path: Path) -> FileInfo:
+        ret: FileInfo | None
+        if not path.exists():
+            ret = FileInfo(path=path, lines=(), code_blocks=[])
+        else:
+            for file in files:
+                if path.samefile(file.path):
+                    return file
+            ret = FileInfo.from_file(file=path)
+            assert ret is not None, f"Failed to read file: {path}"
+        files.append(ret)
+        return ret
+
+    # Step 0. Find out functions and classes already defined on files.
+    defined_func_prefixes: set[str] = {  # type: ignore[union-attr]
+        code.param[0] for file in files for code in file.code_blocks if 
code.kind == "global"
+    }
+    defined_objs: set[str] = {  # type: ignore[assignment]
+        code.param for file in files for code in file.code_blocks if code.kind 
== "object"
+    } | C.BUILTIN_TYPE_KEYS

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   These set comprehensions are quite dense and span multiple lines, which can 
make them difficult to read and understand at a glance. For improved 
readability and maintainability, consider refactoring them into regular `for` 
loops. This would make the logic of iterating through files and code blocks to 
build the sets more explicit and easier to follow.
   
   ```suggestion
       defined_func_prefixes: set[str] = set()
       defined_objs: set[str] = set()
       for file in files:
           for code in file.code_blocks:
               if code.kind == "global":
                   assert isinstance(code.param, tuple)
                   defined_func_prefixes.add(code.param[0])
               elif code.kind == "object":
                   assert isinstance(code.param, str)
                   defined_objs.add(code.param)
       defined_objs.update(C.BUILTIN_TYPE_KEYS)
   ```



-- 
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]

Reply via email to