Merged return tuple to EvalResult class

4 files changed, 53 insertions, 24 deletions

diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py
index 80c9ec4..1e1cc8b 100644
--- a/snekbox/api/resources/eval.py
+++ b/snekbox/api/resources/eval.py
@@ -77,7 +77,7 @@ class EvalResource:
         args = req.media.get("args", ("-c",))
 
         try:
-            result, attachments = self.nsjail.python3(code, py_args=args)
+            result = self.nsjail.python3(code, py_args=args)
         except Exception:
             log.exception("An exception occurred while trying to process the request")
             raise falcon.HTTPInternalServerError
@@ -85,5 +85,5 @@ class EvalResource:
         resp.media = {
             "stdout": result.stdout,
             "returncode": result.returncode,
-            "attachments": [atc.to_dict() for atc in attachments],
+            "attachments": [atc.to_dict() for atc in result.attachments],
         }
diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index 520954b..2eb83d2 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 
 import logging
 import subprocess
+from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cache
 from pathlib import Path
@@ -93,13 +94,12 @@ class MemoryTempDir:
         yield
         self.path.chmod(0o555)
 
-    def attachments(self) -> list[FileAttachment] | None:
+    def attachments(self) -> Generator[FileAttachment, None, None]:
         """Return a list of attachments in the tempdir."""
-        # First look for any file named `output` (any extension)
-        output = next((f for f in self.home.glob("output*") if f.is_file()), None)
-        if output:
-            return [FileAttachment.from_path(output, MAX_FILE_SIZE)]
-        return None
+        # Look for any file starting with `output`
+        for file in self.path.glob("output*"):
+            if file.is_file():
+                yield FileAttachment.from_path(file, MAX_FILE_SIZE)
 
     def cleanup(self) -> None:
         """Remove files in temp dir, releases name."""
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index d8b642a..0344c3c 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -3,19 +3,20 @@ import re
 import subprocess
 import sys
 import textwrap
-from subprocess import CompletedProcess
 from tempfile import NamedTemporaryFile
 from typing import Iterable
 
 from google.protobuf import text_format
 
+# noinspection PyProtectedMember
 from snekbox import DEBUG, utils
 from snekbox.config_pb2 import NsJailConfig
 from snekbox.memfs import MemoryTempDir
 
 __all__ = ("NsJail",)
 
-from snekbox.snekio import FileAttachment
+from snekbox.process import EvalResult
+from snekbox.snekio import AttachmentError
 
 log = logging.getLogger(__name__)
 
@@ -133,7 +134,7 @@ class NsJail:
 
     def python3(
         self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("",)
-    ) -> tuple[CompletedProcess, list[FileAttachment]]:
+    ) -> EvalResult:
         """
         Execute Python 3 code in an isolated environment and return the completed process.
 
@@ -197,20 +198,12 @@ class NsJail:
                     args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
                 )
             except ValueError:
-                return CompletedProcess(args, None, "ValueError: embedded null byte", None), []
+                return EvalResult(args, None, "ValueError: embedded null byte")
 
             try:
                 output = self._consume_stdout(nsjail)
             except UnicodeDecodeError:
-                return (
-                    CompletedProcess(
-                        args,
-                        None,
-                        "UnicodeDecodeError: invalid Unicode in output pipe",
-                        None,
-                    ),
-                    [],
-                )
+                return EvalResult(args, None, "UnicodeDecodeError: invalid Unicode in output pipe")
 
             # When you send signal `N` to a subprocess to terminate it using Popen, it
             # will return `-N` as its exit code. As we normally get `N + 128` back, we
@@ -218,8 +211,12 @@ class NsJail:
             returncode = -nsjail.returncode + 128 if nsjail.returncode < 0 else nsjail.returncode
 
             # Parse attachments
-            attachments = temp_dir.attachments() or []
-            log.info(f"Found {len(attachments)} attachments.")
+            try:
+                attachments = list(temp_dir.attachments())
+                log.info(f"Found {len(attachments)} attachments.")
+            except AttachmentError as err:
+                log.error(f"Failed to parse attachments: {err}")
+                return EvalResult(args, returncode, f"AttachmentError: {err}")
 
             log_lines = nsj_log.read().decode("utf-8").splitlines()
             if not log_lines and returncode == 255:
@@ -230,4 +227,4 @@ class NsJail:
 
         log.info(f"nsjail return code: {returncode}")
 
-        return CompletedProcess(args, returncode, output, None), attachments
+        return EvalResult(args, returncode, output, attachments=attachments)
diff --git a/snekbox/process.py b/snekbox/process.py
new file mode 100644
index 0000000..d6d25c0
--- /dev/null
+++ b/snekbox/process.py
@@ -0,0 +1,32 @@
+"""Utilities for process management."""
+from collections.abc import Sequence
+from os import PathLike
+from subprocess import CompletedProcess
+from typing import TypeVar
+
+from snekbox.snekio import FileAttachment
+
+_T = TypeVar("_T")
+ArgType = (
+    str
+    | bytes
+    | PathLike[str]
+    | PathLike[bytes]
+    | Sequence[str | bytes | PathLike[str] | PathLike[bytes]]
+)
+
+
+class EvalResult(CompletedProcess[_T]):
+    """An evaluation job that has finished running."""
+
+    def __init__(
+        self,
+        args: ArgType,
+        returncode: int | None,
+        stdout: _T | None = None,
+        stderr: _T | None = None,
+        attachments: list[FileAttachment] | None = None,
+    ) -> None:
+        """Create an evaluation result."""
+        super().__init__(args, returncode, stdout, stderr)
+        self.attachments: list[FileAttachment] = attachments or []