The New Computer Update

interpreter/core/computer/computer.py

@@ -1,50 +1,47 @@
-from .languages.applescript import AppleScript
-from .languages.html import HTML
-from .languages.javascript import JavaScript
-from .languages.powershell import PowerShell
-from .languages.python import Python
-from .languages.r import R
-from .languages.shell import Shell
+from .terminal.terminal import Terminal
 
-language_map = {
-    "python": Python,
-    "bash": Shell,
-    "shell": Shell,
-    "sh": Shell,
-    "zsh": Shell,
-    "javascript": JavaScript,
-    "html": HTML,
-    "applescript": AppleScript,
-    "r": R,
-    "powershell": PowerShell,
-}
+try:
+    from .display.display import Display
+    from .keyboard.keyboard import Keyboard
+    from .mouse.mouse import Mouse
+except ImportError or ModuleNotFoundError:
+    raise
+    pass
 
 
 class Computer:
     def __init__(self):
-        self.languages = [Python, Shell, JavaScript, HTML, AppleScript, R, PowerShell]
-        self._active_languages = {}
-
-    def run(self, language, code):
-        if language not in self._active_languages:
-            self._active_languages[language] = language_map[language]()
+        self.terminal = Terminal()
         try:
-            yield from self._active_languages[language].run(code)
-        except GeneratorExit:
-            self.stop()
+            self.mouse = Mouse(
+                self
+            )  # Mouse will use the computer's display, so we give it a reference to ourselves
+            self.keyboard = Keyboard()
+            self.display = Display()
+        except:
+            raise
+            pass
+
+    def run(self, *args, **kwargs):
+        """
+        Shortcut for computer.terminal.run
+        """
+        return self.terminal.run(*args, **kwargs)
 
     def stop(self):
-        for language in self._active_languages.values():
-            language.stop()
+        """
+        Shortcut for computer.terminal.stop
+        """
+        return self.terminal.stop()
 
     def terminate(self):
-        for language_name in list(self._active_languages.keys()):
-            language = self._active_languages[language_name]
-            if (
-                language
-            ):  # Not sure why this is None sometimes. We should look into this
-                language.terminate()
-            del self._active_languages[language_name]
-
-
-computer = Computer()
+        """
+        Shortcut for computer.terminal.terminate
+        """
+        return self.terminal.terminate()
+
+    def screenshot(self, *args, **kwargs):
+        """
+        Shortcut for computer.display.screenshot
+        """
+        return self.display.screenshot(*args, **kwargs)

interpreter/core/computer/display/display.py

@@ -0,0 +1,53 @@
+import os
+import tempfile
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pyautogui
+from PIL import Image
+
+
+class Display:
+    def screenshot(self, show=True, quadrant=None):
+        temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+
+        if quadrant == None:
+            screenshot = pyautogui.screenshot()
+        else:
+            try:
+                screen_width, screen_height = pyautogui.size()
+            except:
+                raise EnvironmentError("Unable to determine screen size.")
+
+            quadrant_width = screen_width // 2
+            quadrant_height = screen_height // 2
+
+            quadrant_coordinates = {
+                1: (0, 0),
+                2: (quadrant_width, 0),
+                3: (0, quadrant_height),
+                4: (quadrant_width, quadrant_height),
+            }
+
+            if quadrant in quadrant_coordinates:
+                x, y = quadrant_coordinates[quadrant]
+                screenshot = pyautogui.screenshot(
+                    region=(x, y, quadrant_width, quadrant_height)
+                )
+            else:
+                raise ValueError("Invalid quadrant. Choose between 1 and 4.")
+
+        screenshot.save(temp_file.name)
+
+        # Open the image file with PIL
+        img = Image.open(temp_file.name)
+
+        # Delete the temporary file
+        os.remove(temp_file.name)
+
+        if show:
+            # Show the image using matplotlib
+            plt.imshow(np.array(img))
+            plt.show()
+
+        return img

interpreter/core/computer/keyboard/keyboard.py

@@ -0,0 +1,38 @@
+import platform
+import random
+import time
+
+import pyautogui
+
+pyautogui.FAILSAFE = False
+
+
+class Keyboard:
+    def write(self, text):
+        # Split the text into words
+        words = text.split(" ")
+
+        # Type each word
+        for word in words:
+            # Type the word
+            pyautogui.write(word)
+            # Add a delay after each word
+            time.sleep(random.uniform(0.1, 0.3))
+
+    def press(self, keys):
+        pyautogui.press(keys)
+
+    def hotkey(self, *args):
+        if "darwin" in platform.system().lower():
+            # For some reason, application focus or something, we need to do this for spotlight
+            # only if they passed in "command", "space" or "command", " ", or those in another order
+            if set(args) == {"command", " "} or set(args) == {"command", "space"}:
+                pyautogui.click(0, 0)
+                time.sleep(0.5)
+        pyautogui.hotkey(*args)
+
+    def down(self, key):
+        pyautogui.keyDown(key)
+
+    def up(self, key):
+        pyautogui.keyUp(key)

interpreter/core/computer/mouse/mouse.py

@@ -0,0 +1,79 @@
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pyautogui
+from PIL import Image
+
+from ..utils.computer_vision import find_text_in_image
+
+pyautogui.FAILSAFE = False
+
+
+class Mouse:
+    def __init__(self, computer):
+        self.computer = computer
+
+    def move(self, *args, x=None, y=None, index=None, svg=None):
+        if len(args) > 1:
+            raise ValueError(
+                "Too many positional arguments provided: click(*args, x=None, y=None, show=True, index=None)"
+            )
+        elif len(args) == 1:
+            text = args[0]
+            # Take a screenshot
+            img = self.computer.screenshot(show=False)
+
+            # Find the text in the screenshot
+            centers, bounding_box_image = find_text_in_image(img, text)
+
+            # If the text was found
+            if centers:
+                # This could be refactored to be more readable
+                if len(centers) > 1:
+                    if index == None:
+                        print(
+                            f"This text ('{text}') was found multiple times on screen. Please try 'click()' again, but pass in an `index` int to identify which one you want to click. The indices have been drawn on the attached image."
+                        )
+                        # Show the image using matplotlib
+                        plt.imshow(np.array(bounding_box_image))
+                        plt.show()
+                        return
+                    else:
+                        center = centers[index]
+                else:
+                    center = centers[0]
+
+                # Slowly move the mouse from its current position
+                pyautogui.moveTo(center[0], center[1], duration=0.5)
+
+            else:
+                plt.imshow(np.array(bounding_box_image))
+                plt.show()
+                print("Your text was not found on the screen. Please try again.")
+        elif x is not None and y is not None:
+            # Move to the specified coordinates and click
+            pyautogui.moveTo(x, y, duration=0.5)
+        elif svg is not None:
+            raise NotImplementedError("SVG handling not implemented yet.")
+            # img = self.computer.screenshot(show=False)
+            # # Move to the specified coordinates and click
+            # coordinates = find_svg_in_image(svg, img)
+            # if coordinates == None:
+            #     print("Not found.")
+            #     return
+            # pyautogui.moveTo(coordinates[0], coordinates[1], duration=0.5)
+            # pyautogui.click(coordinates[0], coordinates[1])
+        else:
+            raise ValueError("Either text or both x and y must be provided")
+
+    def click(self, *args, **kwargs):
+        if args or kwargs:
+            self.move(*args, **kwargs)
+        pyautogui.click()
+
+    def down(self):
+        pyautogui.mouseDown()
+
+    def up(self):
+        pyautogui.mouseUp()

interpreter/core/computer/languages/ARCHIVE_subprocess_python.py → interpreter/core/computer/terminal/languages/ARCHIVE_subprocess_python.py

@@ -6,7 +6,7 @@
 import shlex
 import sys
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class Python(SubprocessLanguage):

interpreter/core/computer/languages/applescript.py → interpreter/core/computer/terminal/languages/applescript.py

@@ -1,6 +1,6 @@
 import os
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class AppleScript(SubprocessLanguage):

interpreter/core/computer/languages/javascript.py → interpreter/core/computer/terminal/languages/javascript.py

@@ -1,6 +1,6 @@
 import re
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class JavaScript(SubprocessLanguage):

interpreter/core/computer/jupyter_language.py → interpreter/core/computer/terminal/languages/jupyter_language.py

@@ -3,11 +3,12 @@
 import re
 import threading
 import time
+import traceback
 
 import matplotlib
 from jupyter_client import KernelManager
 
-from .base_language import BaseLanguage
+from ..base_language import BaseLanguage
 
 DEBUG_MODE = False
 
@@ -47,10 +48,16 @@ def terminate(self):
 
     def run(self, code):
         self.finish_flag = False
-        preprocessed_code = self.preprocess_code(code)
-        message_queue = queue.Queue()
-        self._execute_code(preprocessed_code, message_queue)
-        yield from self._capture_output(message_queue)
+        try:
+            preprocessed_code = self.preprocess_code(code)
+            message_queue = queue.Queue()
+            self._execute_code(preprocessed_code, message_queue)
+            yield from self._capture_output(message_queue)
+        except GeneratorExit:
+            raise  # gotta pass this up!
+        except:
+            content = traceback.format_exc()
+            yield {"type": "console", "format": "output", "content": content}
 
     def _execute_code(self, code, message_queue):
         def iopub_message_listener():
@@ -197,6 +204,8 @@ def preprocess_python(code):
     Wrap in a try except
     """
 
+    code = code.strip()
+
     # Add print commands that tell us what the active line is
     code = add_active_line_prints(code)
 
@@ -216,6 +225,13 @@ def add_active_line_prints(code):
     """
     Add print statements indicating line numbers to a python string.
     """
+    # Replace newlines and comments with pass statements, so the line numbers are accurate (ast will remove them otherwise)
+    code_lines = code.split("\n")
+    for i in range(len(code_lines)):
+        if code_lines[i].strip().startswith("#") or code_lines[i] == "":
+            whitespace = len(code_lines[i]) - len(code_lines[i].lstrip())
+            code_lines[i] = " " * whitespace + "pass"
+    code = "\n".join(code_lines)
     tree = ast.parse(code)
     transformer = AddLinePrints()
     new_tree = transformer.visit(tree)

interpreter/core/computer/languages/powershell.py → interpreter/core/computer/terminal/languages/powershell.py

@@ -2,7 +2,7 @@
 import platform
 import shutil
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class PowerShell(SubprocessLanguage):

interpreter/core/computer/languages/python.py → interpreter/core/computer/terminal/languages/python.py

@@ -1,6 +1,6 @@
 import os
 
-from ..jupyter_language import JupyterLanguage
+from .jupyter_language import JupyterLanguage
 
 # Supresses a weird debugging error
 os.environ["PYDEVD_DISABLE_FILE_VALIDATION"] = "1"

interpreter/core/computer/languages/r.py → interpreter/core/computer/terminal/languages/r.py

@@ -1,6 +1,6 @@
 import re
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class R(SubprocessLanguage):

interpreter/core/computer/languages/shell.py → interpreter/core/computer/terminal/languages/shell.py

@@ -2,7 +2,7 @@
 import platform
 import re
 
-from ..subprocess_language import SubprocessLanguage
+from .subprocess_language import SubprocessLanguage
 
 
 class Shell(SubprocessLanguage):

interpreter/core/computer/subprocess_language.py → interpreter/core/computer/terminal/languages/subprocess_language.py

@@ -6,7 +6,7 @@
 import time
 import traceback
 
-from .base_language import BaseLanguage
+from ..base_language import BaseLanguage
 
 
 class SubprocessLanguage(BaseLanguage):