Insert Python functions into the TVMScript output.
(self, base_script: str, indent_spaces: int)
| 544 | return self._insert_python_functions(base_script, indent_spaces) |
| 545 | |
| 546 | def _insert_python_functions(self, base_script: str, indent_spaces: int) -> str: |
| 547 | """Insert Python functions into the TVMScript output.""" |
| 548 | lines = base_script.split("\n") |
| 549 | result_lines = [] |
| 550 | |
| 551 | # Find the class definition line and insert Python functions after it |
| 552 | class_found = False |
| 553 | class_indent = 0 |
| 554 | |
| 555 | for line in lines: |
| 556 | result_lines.append(line) |
| 557 | |
| 558 | # Look for class definition |
| 559 | if not class_found and line.strip().startswith("class "): |
| 560 | class_found = True |
| 561 | class_indent = len(line) - len(line.lstrip()) |
| 562 | |
| 563 | # Insert Python functions after the class definition |
| 564 | if hasattr(self.ir_mod, "pyfuncs") and self.ir_mod.pyfuncs: |
| 565 | for func_name, func in self.ir_mod.pyfuncs.items(): |
| 566 | # Get the function source code |
| 567 | func_source = self._get_function_source(func) |
| 568 | if func_source: |
| 569 | # Format the function with proper indentation |
| 570 | formatted_func = self._format_python_function( |
| 571 | func_name, func_source, class_indent + indent_spaces |
| 572 | ) |
| 573 | result_lines.append(formatted_func) |
| 574 | result_lines.append("") # Add empty line for separation |
| 575 | |
| 576 | return "\n".join(result_lines) |
| 577 | |
| 578 | def _get_function_source(self, func: callable) -> str | None: |
| 579 | """Get the source code of a Python function.""" |
no test coverage detected