From fe910040610099a91ebb1cb1008d7206648717b1 Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Thu, 14 Nov 2024 12:19:30 -0800 Subject: [PATCH] add tasks with ast --- agentstack/generation/gen_utils.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/agentstack/generation/gen_utils.py b/agentstack/generation/gen_utils.py index 25b2889..3faf78c 100644 --- a/agentstack/generation/gen_utils.py +++ b/agentstack/generation/gen_utils.py @@ -27,17 +27,31 @@ def insert_after_tasks(file_path, code_to_insert): module = ast.parse(content) - # Track the last task function's end line + # Track the last task function and its line number last_task_end = None + last_task_start = None for node in ast.walk(module): if isinstance(node, ast.FunctionDef) and \ any(isinstance(deco, ast.Name) and deco.id == 'task' for deco in node.decorator_list): last_task_end = node.end_lineno + last_task_start = node.lineno if last_task_end is not None: lines = content.split('\n') - for i, line in enumerate(code_to_insert): - lines.insert(last_task_end + i, line) + + # Get the indentation of the task function + task_line = lines[last_task_start - 1] # -1 for 0-based indexing + indentation = '' + for char in task_line: + if char in [' ', '\t']: + indentation += char + else: + break + + # Add the same indentation to each line of the inserted code + indented_code = '\n' + '\n'.join(indentation + line for line in code_to_insert) + + lines.insert(last_task_end, indented_code) content = '\n'.join(lines) with open(file_path, 'w') as file: