From 9506b46c60e4b71957ba2b9ebaffe3772c2c84b6 Mon Sep 17 00:00:00 2001 From: Guilherme Ramos Date: Thu, 14 Nov 2024 22:05:28 -0500 Subject: [PATCH] Refactor Shell class to enhance subprocess handling and improve active line detection in shell scripts. --- .../core/computer/terminal/languages/shell.py | 205 +++++++++++++----- .../terminal/languages/subprocess_language.py | 6 + 2 files changed, 155 insertions(+), 56 deletions(-) diff --git a/interpreter/core/computer/terminal/languages/shell.py b/interpreter/core/computer/terminal/languages/shell.py index 6b900b47e2..db8eedcdc3 100644 --- a/interpreter/core/computer/terminal/languages/shell.py +++ b/interpreter/core/computer/terminal/languages/shell.py @@ -1,41 +1,12 @@ import os -import platform +import queue import re +import subprocess +import threading from .subprocess_language import SubprocessLanguage -class Shell(SubprocessLanguage): - file_extension = "sh" - name = "Shell" - aliases = ["bash", "sh", "zsh", "batch", "bat"] - - def __init__( - self, - ): - super().__init__() - - # Determine the start command based on the platform - if platform.system() == "Windows": - self.start_cmd = ["cmd.exe"] - else: - self.start_cmd = [os.environ.get("SHELL", "bash")] - - def preprocess_code(self, code): - return preprocess_shell(code) - - def line_postprocessor(self, line): - return line - - def detect_active_line(self, line): - if "##active_line" in line: - return int(line.split("##active_line")[1].split("##")[0]) - return None - - def detect_end_of_execution(self, line): - return "##end_of_execution##" in line - - def preprocess_shell(code): """ Add active line markers @@ -46,9 +17,9 @@ def preprocess_shell(code): # Add commands that tell us what the active line is # if it's multiline, just skip this. soon we should make it work with multiline if ( - not has_multiline_commands(code) - and os.environ.get("INTERPRETER_ACTIVE_LINE_DETECTION", "True").lower() - == "true" + not has_multiline_commands(code) + and os.environ.get("INTERPRETER_ACTIVE_LINE_DETECTION", "True").lower() + == "true" ): code = add_active_line_prints(code) @@ -57,6 +28,32 @@ def preprocess_shell(code): return code +def has_multiline_commands(script_text): + """ + Check if the shell script contains multiline commands. + """ + # Patterns that indicate a line continues + continuation_patterns = [ + r"\\$", # Line continuation character at the end of the line + r"\|$", # Pipe character at the end of the line indicating a pipeline continuation + r"&&\s*$", # Logical AND at the end of the line + r"\|\|\s*$", # Logical OR at the end of the line + r"<\($", # Start of process substitution + r"\($", # Start of subshell + r"{\s*$", # Start of a block + r"\bif\b", # Start of an if statement + r"\bwhile\b", # Start of a while loop + r"\bfor\b", # Start of a for loop + r"do\s*$", # 'do' keyword for loops + r"then\s*$", # 'then' keyword for if statements + ] + + # Check each line for multiline patterns + for line in script_text.splitlines(): + if any(re.search(pattern, line.rstrip()) for pattern in continuation_patterns): + return True + + return False def add_active_line_prints(code): """ @@ -64,31 +61,127 @@ def add_active_line_prints(code): """ lines = code.split("\n") for index, line in enumerate(lines): + # Skip empty lines + if not line.strip(): + continue # Insert the echo command before the actual line lines[index] = f'echo "##active_line{index + 1}##"\n{line}' return "\n".join(lines) +def preprocess_shell(code): + # Add commands that tell us what the active line is + if ( + not has_multiline_commands(code) + and os.environ.get("INTERPRETER_ACTIVE_LINE_DETECTION", "True").lower() + == "true" + ): + code = add_active_line_prints(code) -def has_multiline_commands(script_text): - # Patterns that indicate a line continues - continuation_patterns = [ - r"\\$", # Line continuation character at the end of the line - r"\|$", # Pipe character at the end of the line indicating a pipeline continuation - r"&&\s*$", # Logical AND at the end of the line - r"\|\|\s*$", # Logical OR at the end of the line - r"<\($", # Start of process substitution - r"\($", # Start of subshell - r"{\s*$", # Start of a block - r"\bif\b", # Start of an if statement - r"\bwhile\b", # Start of a while loop - r"\bfor\b", # Start of a for loop - r"do\s*$", # 'do' keyword for loops - r"then\s*$", # 'then' keyword for if statements - ] + # Add end command (we'll be listening for this so we know when it ends) + code += '\necho "##end_of_execution##"' + + return code + +class Shell(SubprocessLanguage): + file_extension = "sh" + name = "Shell" + aliases = ["bash", "sh", "zsh", "batch", "bat"] + + def run(self, code): + code = preprocess_shell(code) + self.output_queue = queue.Queue() + self.done = threading.Event() + + process = subprocess.Popen( + code, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, # Prevent subprocess from waiting for input + text=True + ) + + threading.Thread( + target=self.handle_stream_output, + args=(process.stdout, False), + daemon=True, + ).start() + threading.Thread( + target=self.handle_stream_output, + args=(process.stderr, True), + daemon=True, + ).start() + + threading.Thread( + target=self.wait_for_process, + args=(process,), + daemon=True, + ).start() + + while True: + try: + output = self.output_queue.get(timeout=0.1) + yield output + except queue.Empty: + if self.done.is_set(): + break + + def handle_stream_output(self, stream, is_error_stream): + try: + for line in iter(stream.readline, ""): + line = line.rstrip('\n') + if self.detect_active_line(line): + active_line = self.detect_active_line(line) + self.output_queue.put( + { + "type": "console", + "format": "active_line", + "content": active_line, + } + ) + line = re.sub(r"##active_line\d+##", "", line) + if line: + self.output_queue.put( + {"type": "console", "format": "output", "content": line} + ) + elif self.detect_end_of_execution(line): + line = line.replace("##end_of_execution##", "").strip() + if line: + self.output_queue.put( + {"type": "console", "format": "output", "content": line} + ) + else: + self.output_queue.put( + {"type": "console", "format": "output", "content": line} + ) + except Exception as e: + self.output_queue.put( + {"type": "console", "format": "error", "content": str(e)} + ) + finally: + stream.close() + + def wait_for_process(self, process, timeout=30): + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + process.kill() + self.output_queue.put( + { + "type": "console", + "format": "error", + "content": f"Process timed out after {timeout} seconds.", + } + ) + finally: + self.done.set() + + def detect_active_line(self, line): + if "##active_line" in line: + return int(line.split("##active_line")[1].split("##")[0]) + return None + + def detect_end_of_execution(self, line): + return "##end_of_execution##" in line - # Check each line for multiline patterns - for line in script_text.splitlines(): - if any(re.search(pattern, line.rstrip()) for pattern in continuation_patterns): - return True - return False diff --git a/interpreter/core/computer/terminal/languages/subprocess_language.py b/interpreter/core/computer/terminal/languages/subprocess_language.py index dd422beb7f..1d698c328b 100644 --- a/interpreter/core/computer/terminal/languages/subprocess_language.py +++ b/interpreter/core/computer/terminal/languages/subprocess_language.py @@ -138,6 +138,7 @@ def run(self, code): break def handle_stream_output(self, stream, is_error_stream): + self.verbose = True try: for line in iter(stream.readline, ""): if self.verbose: @@ -191,3 +192,8 @@ def handle_stream_output(self, stream, is_error_stream): print("Stream closed while reading.") else: raise e + finally: + # Set self.done when the stream ends + if self.verbose: + print("Stream ended, setting self.done") + self.done.set()