Run plugin: add memory limit

porcupine/plugins/run/__init__.py

@@ -12,6 +12,7 @@
 
 from porcupine import get_main_window, get_tab_manager, menubar, tabs, utils
 from porcupine.plugins import python_venv
+from porcupine.settings import global_settings
 
 from . import common, dialog, history, no_terminal, terminal
 
@@ -76,6 +77,11 @@ def on_new_filetab(tab: tabs.FileTab) -> None:
 
 
 def setup() -> None:
+    # Memory limit affects all commands, so that you can "set it and forget it".
+    # If you choose to run a different command, it won't reset annoyingly.
+    global_settings.add_option("run_mem_limit_enabled", default=False)
+    global_settings.add_option("run_mem_limit_value", default=1000 * 1000 * 1000)
+
     get_tab_manager().add_filetab_callback(on_new_filetab)
 
     menubar.add_filetab_command(

porcupine/plugins/run/common.py

@@ -2,10 +2,13 @@
 
 import dataclasses
 import os
+import resource
+import sys
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Callable, Dict, List, Optional
 
 from porcupine import tabs, utils
+from porcupine.settings import global_settings
 
 
 @dataclasses.dataclass
@@ -34,7 +37,7 @@ class ExampleCommand:
 
 
 class Context:
-    def __init__(self, tab: tabs.FileTab, key_id: int):
+    def __init__(self, tab: tabs.FileTab, key_id: int) -> None:
         assert tab.path is not None
         self.file_path = tab.path
         self.project_path = utils.find_project_root(tab.path)
@@ -69,3 +72,13 @@ def prepare_env() -> dict[str, str]:
         )
 
     return env
+
+
+def create_memory_limit_callback() -> Callable[[], None]:
+    assert sys.platform != "win32"
+
+    if global_settings.get("run_mem_limit_enabled", bool):
+        limit = global_settings.get("run_mem_limit_value", int)
+        return lambda: resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
+    else:
+        return lambda: None

porcupine/plugins/run/dialog.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import re
 import sys
 import tkinter
 from tkinter import ttk
 from typing import Callable, TypeVar
 
 from porcupine import get_main_window, textutils, utils
+from porcupine.settings import global_settings
 
 from . import common, history
 
@@ -56,6 +58,51 @@ def validate(self, *junk_from_var_trace: object) -> None:
         self._validated_callback()
 
 
+def _mem_limit_to_string(limit: int) -> str:
+    if limit >= 1000 * 1000 * 1000:
+        number = limit / (1000 * 1000 * 1000)
+        suffix = "GB"
+    elif limit >= 1000 * 1000:
+        number = limit / (1000 * 1000)
+        suffix = "MB"
+    elif limit >= 1000:
+        number = limit / 1000
+        suffix = "KB"
+    else:
+        number = limit
+        suffix = "B"
+
+    # Show 2GB instead of 2.0GB
+    if number == int(number):
+        number = int(number)
+
+    return str(number) + suffix
+
+
+def _string_to_mem_limit(string: str) -> int | None:
+    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)([KMG]?)B?", string)
+    if match is None:
+        return None
+
+    number_as_string, suffix = match.groups()
+    number = float(number_as_string)
+    if suffix == "":
+        limit = round(number)
+    elif suffix == "K":
+        limit = round(1000 * number)
+    elif suffix == "M":
+        limit = round(1000 * 1000 * number)
+    elif suffix == "G":
+        limit = round(1000 * 1000 * 1000 * number)
+    else:
+        raise NotImplementedError
+
+    # Ban limits smaller than 10MB. Python needs about 17MB to start.
+    if limit < 10 * 1000 * 1000:
+        return None
+    return limit
+
+
 class _CommandAsker:
     def __init__(self, ctx: common.Context):
         self.window = tkinter.Toplevel(name="run_command_asker")
@@ -77,13 +124,13 @@ def __init__(self, ctx: common.Context):
             entry_area,
             text="Run this command:",
             get_substituted_value=(lambda: self.get_command().format_command()),
-            validated_callback=self.update_run_button,
+            validated_callback=self.update_disabled_states,
         )
         self.cwd = _FormattingEntryAndLabels(
             entry_area,
             text="In this directory:",
             get_substituted_value=(lambda: str(self.get_command().format_cwd())),
-            validated_callback=self.update_run_button,
+            validated_callback=self.update_disabled_states,
         )
 
         ttk.Label(content_frame, text="Substitutions:").pack(anchor="w")
@@ -120,11 +167,30 @@ def __init__(self, ctx: common.Context):
         self.window.bind("<Alt-p>", (lambda e: self.terminal_var.set(False)), add=True)
         self.window.bind("<Alt-e>", (lambda e: self.terminal_var.set(True)), add=True)
 
+        self._mem_limit_enable_var = tkinter.BooleanVar(
+            value=global_settings.get("run_mem_limit_enabled", bool)
+        )
+        self._mem_limit_value_var = tkinter.StringVar(
+            value=_mem_limit_to_string(global_settings.get("run_mem_limit_value", int))
+        )
+
+        memory_frame = ttk.Frame(content_frame)
+        if sys.platform != "win32":
+            memory_frame.pack(fill="x", pady=5)
+        ttk.Checkbutton(
+            memory_frame, variable=self._mem_limit_enable_var, text="Limit memory usage to: "
+        ).pack(side="left")
+        self._mem_entry = ttk.Entry(memory_frame, textvariable=self._mem_limit_value_var, width=20)
+        self._mem_entry.pack(side="left")
+
+        self._mem_limit_enable_var.trace_add("write", self.update_disabled_states)
+        self._mem_limit_value_var.trace_add("write", self.update_disabled_states)
+
         self._repeat_bindings = [
             utils.get_binding(f"<<Run:Repeat{key_id}>>") for key_id in range(4)
         ]
         self._repeat_var = tkinter.StringVar(value=self._repeat_bindings[ctx.key_id])
-        self._repeat_var.trace_add("write", self.update_run_button)
+        self._repeat_var.trace_add("write", self.update_disabled_states)
 
         repeat_frame = ttk.Frame(content_frame)
         repeat_frame.pack(fill="x", pady=10)
@@ -145,7 +211,7 @@ def __init__(self, ctx: common.Context):
         self.run_button.pack(side="left", fill="x", expand=True, padx=(5, 0))
         self.run_clicked = False
 
-        for entry in [self.command.entry, self.cwd.entry]:
+        for entry in [self.command.entry, self.cwd.entry, self._mem_entry]:
             entry.bind("<Return>", (lambda e: self.run_button.invoke()), add=True)
             entry.bind("<Escape>", (lambda e: self.window.destroy()), add=True)
 
@@ -174,16 +240,32 @@ def get_command(self) -> common.Command:
             substitutions=self._substitutions,
         )
 
+    def save_memory_limit_setting(self) -> None:
+        global_settings.set("run_mem_limit_enabled", self._mem_limit_enable_var.get())
+        value = _string_to_mem_limit(self._mem_limit_value_var.get())
+        if value is not None:
+            global_settings.set("run_mem_limit_value", value)
+
     def get_key_id(self) -> int:
         return self._repeat_bindings.index(self._repeat_var.get())
 
-    def command_and_cwd_are_valid(self) -> bool:
+    def _all_entries_are_valid(self) -> bool:
         try:
             command = self.get_command().format_command()
             cwd = self.get_command().format_cwd()
         except (ValueError, KeyError, IndexError):
             return False
-        return bool(command.strip()) and cwd.is_dir() and cwd.is_absolute()
+
+        return (
+            bool(command.strip())
+            and cwd.is_dir()
+            and cwd.is_absolute()
+            and self._repeat_var.get() in self._repeat_bindings
+            and (
+                _string_to_mem_limit(self._mem_limit_value_var.get()) is not None
+                or not self._mem_limit_enable_var.get()
+            )
+        )
 
     def _select_command_autocompletion(self, command: common.Command, prefix: str) -> None:
         assert command.command_format.startswith(prefix)
@@ -210,12 +292,17 @@ def _autocomplete(self, event: tkinter.Event[tkinter.Entry]) -> str | None:
 
         return None
 
-    def update_run_button(self, *junk: object) -> None:
-        if self.command_and_cwd_are_valid() and self._repeat_var.get() in self._repeat_bindings:
+    def update_disabled_states(self, *junk: object) -> None:
+        if self._all_entries_are_valid():
             self.run_button.config(state="normal")
         else:
             self.run_button.config(state="disabled")
 
+        if self._mem_limit_enable_var.get():
+            self._mem_entry.config(state="normal")
+        else:
+            self._mem_entry.config(state="disabled")
+
     def on_run_clicked(self) -> None:
         self.run_clicked = True
         self.window.destroy()
@@ -236,5 +323,6 @@ def ask_command(ctx: common.Context) -> tuple[common.Command, int] | None:
     asker.window.wait_window()
 
     if asker.run_clicked:
+        asker.save_memory_limit_setting()
         return (asker.get_command(), asker.get_key_id())
     return None

porcupine/plugins/run/history.py

@@ -3,6 +3,7 @@
 import copy
 import dataclasses
 import json
+import logging
 import sys
 from pathlib import Path
 from typing import Optional
@@ -13,6 +14,8 @@
 
 from . import common
 
+_log = logging.getLogger(__name__)
+
 
 @dataclasses.dataclass
 class _HistoryItem:

porcupine/plugins/run/no_terminal.py

@@ -98,6 +98,10 @@ def running(self) -> bool:
     def _thread_target(self, command: str, env: dict[str, str]) -> None:
         self._queue.put(("info", command + "\n"))
 
+        popen_kwargs = utils.subprocess_kwargs.copy()
+        if sys.platform != "win32":
+            popen_kwargs["preexec_fn"] = common.create_memory_limit_callback()
+
         try:
             self._shell_process = subprocess.Popen(
                 command,
@@ -106,7 +110,7 @@ def _thread_target(self, command: str, env: dict[str, str]) -> None:
                 stderr=subprocess.STDOUT,
                 shell=True,
                 env=env,
-                **utils.subprocess_kwargs,
+                **popen_kwargs,
             )
         except OSError as e:
             self._queue.put(("error", f"{type(e).__name__}: {e}\n"))

porcupine/plugins/run/terminal.py

@@ -34,6 +34,8 @@ def _run_in_windows_cmd(command: str, cwd: Path, env: dict[str, str]) -> None:
         # new command prompt i found and it works :) we need cmd
         # because start is built in to cmd (lol)
         real_command = ["cmd", "/c", "start"] + real_command
+
+    # no memory limit, don't know how to get that to work on windows
     subprocess.Popen(real_command, env=env)
 
 
@@ -53,7 +55,11 @@ def _run_in_macos_terminal_app(command: str, cwd: Path, env: dict[str, str]) ->
         )
 
     os.chmod(file.name, 0o755)
-    subprocess.Popen(["open", "-a", "Terminal.app", file.name], env=env)
+    subprocess.Popen(
+        ["open", "-a", "Terminal.app", file.name],
+        env=env,
+        preexec_fn=common.create_memory_limit_callback(),
+    )
     # the terminal might be still opening when we get here, that's why
     # the file deletes itself
     # right now the file removes itself before it runs the actual command so
@@ -118,7 +124,11 @@ def _run_in_x11_like_terminal(command: str, cwd: Path, env: dict[str, str]) -> N
 
     real_command = [str(run_script), str(cwd), command]
     real_command.extend(map(str, command))
-    subprocess.Popen([terminal, "-e", " ".join(map(shlex.quote, real_command))], env=env)
+    subprocess.Popen(
+        [terminal, "-e", " ".join(map(shlex.quote, real_command))],
+        env=env,
+        preexec_fn=common.create_memory_limit_callback(),
+    )
 
 
 # this figures out which terminal to use every time the user wants to run

tests/test_run_plugin.py

@@ -1,5 +1,6 @@
 import os
 import shutil
+import struct
 import sys
 import time
 from tkinter import ttk
@@ -8,6 +9,7 @@
 
 from porcupine import get_main_window, get_tab_manager, utils
 from porcupine.plugins.run import common, dialog, history, no_terminal, terminal
+from porcupine.settings import global_settings
 
 
 @pytest.fixture(autouse=True)
@@ -310,6 +312,34 @@ def test_crlf_on_any_platform(tmp_path, wait_until):
     wait_until(lambda: "foo\nbar" in get_output())
 
 
+@pytest.fixture(autouse=True)
+def restore_memory_limits():
+    yield
+    global_settings.reset("run_mem_limit_enabled")
+    global_settings.reset("run_mem_limit_value")
+
+
+@pytest.mark.skipif(struct.calcsize("P") != 8, reason="assumes 64-bit")
+@pytest.mark.xfail(
+    sys.platform == "win32", strict=True, reason="memory limits don't work on windows"
+)
+def test_memory_limit(tmp_path, wait_until):
+    global_settings.set("run_mem_limit_enabled", True)
+    global_settings.set("run_mem_limit_value", 100 * 1000 * 1000)  # 100MB
+
+    # 1 array element = pointer = 8 bytes (on a 64-bit system)
+    # 15*1000*1000 array elements = 120MB
+    no_terminal.run_command(f'{utils.quote(sys.executable)} -c "[0] * (15*1000*1000)"', tmp_path)
+    wait_until(lambda: "The process failed" in get_output())
+    assert "MemoryError" in get_output()
+
+    global_settings.set("run_mem_limit_value", 1000 * 1000 * 1000)  # 1GB
+
+    no_terminal.run_command(f'{utils.quote(sys.executable)} -c "[0] * (15*1000*1000)"', tmp_path)
+    wait_until(lambda: "The process completed successfully" in get_output())
+    assert "MemoryError" not in get_output()
+
+
 def test_changing_current_file(filetab, tmp_path, wait_until):
     filetab.textwidget.insert("end", 'with open("foo.py", "w") as f: f.write("lol")')
     filetab.save_as(tmp_path / "foo.py")