From 64ca95ceeb6f3c919ef7c1c607369e79e32b663c Mon Sep 17 00:00:00 2001
From: Mark <1515135+MarkKoz@users.noreply.github.com>
Date: Sat, 16 Sep 2023 12:46:44 -0700
Subject: Refactor writing of input files into a separate function

---
 snekbox/nsjail.py     | 39 +++++++++++++++++++++++----------------
 snekbox/process.py    | 32 --------------------------------
 snekbox/result.py     | 38 ++++++++++++++++++++++++++++++++++++++
 tests/api/__init__.py |  2 +-
 4 files changed, 62 insertions(+), 49 deletions(-)
 delete mode 100644 snekbox/process.py
 create mode 100644 snekbox/result.py

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 4e86f33..b1d6547 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -12,7 +12,7 @@ from google.protobuf import text_format
 from snekbox import DEBUG, limits
 from snekbox.config_pb2 import NsJailConfig
 from snekbox.limits.timed import time_limit
-from snekbox.process import EvalResult
+from snekbox.result import EvalError, EvalResult
 from snekbox.snekio import FileAttachment, MemFS
 from snekbox.snekio.filesystem import Size
 from snekbox.utils.iter import iter_lstrip
@@ -199,6 +199,24 @@ class NsJail:
             *iter_lstrip(py_args),
         ]
 
+    def _write_files(self, home: Path, files: Iterable[FileAttachment]) -> dict[Path, float]:
+        files_written = {}
+        for file in files:
+            try:
+                f_path = file.save_to(home)
+                # Allow file to be writable
+                f_path.chmod(0o777)
+                # Save the written at time to later check if it was modified
+                files_written[f_path] = f_path.stat().st_mtime
+                log.info(f"Created file at {(home / file.path)!r}.")
+            except OSError as e:
+                log.info(f"Failed to create file at {(home / file.path)!r}.", exc_info=e)
+                raise EvalError(
+                    f"{e.__class__.__name__}: Failed to create file '{file.path}'."
+                ) from e
+
+        return files_written
+
     def python3(
         self,
         py_args: Iterable[str],
@@ -220,21 +238,10 @@ class NsJail:
         ) as fs:
             args = self._build_args(py_args, nsjail_args, nsj_log.name, str(fs.home))
 
-            # Write provided files if any
-            files_written: dict[Path, float] = {}
-            for file in files:
-                try:
-                    f_path = file.save_to(fs.home)
-                    # Allow file to be writable
-                    f_path.chmod(0o777)
-                    # Save the written at time to later check if it was modified
-                    files_written[f_path] = f_path.stat().st_mtime
-                    log.info(f"Created file at {(fs.home / file.path)!r}.")
-                except OSError as e:
-                    log.info(f"Failed to create file at {(fs.home / file.path)!r}.", exc_info=e)
-                    return EvalResult(
-                        args, None, f"{e.__class__.__name__}: Failed to create file '{file.path}'."
-                    )
+            try:
+                files_written = self._write_files(fs.home, files)
+            except EvalError as e:
+                return EvalResult(args, None, str(e))
 
             msg = "Executing code..."
             if DEBUG:
diff --git a/snekbox/process.py b/snekbox/process.py
deleted file mode 100644
index 552b91a..0000000
--- a/snekbox/process.py
+++ /dev/null
@@ -1,32 +0,0 @@
-"""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,
-        files: list[FileAttachment] | None = None,
-    ) -> None:
-        """Create an evaluation result."""
-        super().__init__(args, returncode, stdout, stderr)
-        self.files: list[FileAttachment] = files or []
diff --git a/snekbox/result.py b/snekbox/result.py
new file mode 100644
index 0000000..e11d0e4
--- /dev/null
+++ b/snekbox/result.py
@@ -0,0 +1,38 @@
+"""Types for representing the result of an evaluation job."""
+from collections.abc import Sequence
+from os import PathLike
+from subprocess import CompletedProcess
+from typing import TypeVar
+
+from snekbox.snekio import FileAttachment
+
+__all__ = ("EvalError", "EvalResult")
+
+_T = TypeVar("_T")
+ArgType = (
+    str
+    | bytes
+    | PathLike[str]
+    | PathLike[bytes]
+    | Sequence[str | bytes | PathLike[str] | PathLike[bytes]]
+)
+
+
+class EvalError(RuntimeError):
+    """An error that occurred during evaluation."""
+
+
+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,
+        files: list[FileAttachment] | None = None,
+    ) -> None:
+        """Create an evaluation result."""
+        super().__init__(args, returncode, stdout, stderr)
+        self.files: list[FileAttachment] = files or []
diff --git a/tests/api/__init__.py b/tests/api/__init__.py
index 5f20faf..c84080e 100644
--- a/tests/api/__init__.py
+++ b/tests/api/__init__.py
@@ -4,7 +4,7 @@ from unittest import mock
 from falcon import testing
 
 from snekbox.api import SnekAPI
-from snekbox.process import EvalResult
+from snekbox.result import EvalResult
 
 
 class SnekAPITestCase(testing.TestCase):
-- 
cgit v1.2.3