Add configurable home and output paths

2 files changed, 27 insertions, 6 deletions

diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index d02c029..259f6f7 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -22,7 +22,13 @@ __all__ = ("MemFS",)
 class MemFS:
     """An in-memory temporary file system."""
 
-    def __init__(self, instance_size: int, root_dir: str | Path = "/memfs") -> None:
+    def __init__(
+        self,
+        instance_size: int,
+        root_dir: str | Path = "/memfs",
+        home: str = "home",
+        output: str = "output",
+    ) -> None:
         """
         Initialize an in-memory temporary file system.
 
@@ -33,10 +39,14 @@ class MemFS:
         Args:
             instance_size: Size limit of each tmpfs instance in bytes.
             root_dir: Root directory to mount instances in.
+            home: Name of the home directory.
+            output: Name of the output directory within home. If empty, uses home.
         """
         self.instance_size = instance_size
         self.root_dir = Path(root_dir)
         self.root_dir.mkdir(exist_ok=True, parents=True)
+        self._home_name = home
+        self._output_name = output
 
         for _ in range(10):
             name = str(uuid4())
@@ -82,12 +92,12 @@ class MemFS:
     @property
     def home(self) -> Path:
         """Path to home directory."""
-        return self.path / "home"
+        return self.path / self._home_name
 
     @property
     def output(self) -> Path:
         """Path to output directory."""
-        return self.home / "output"
+        return self.home / self._output_name
 
     def __enter__(self) -> MemFS:
         return self
@@ -112,7 +122,7 @@ class MemFS:
 
     def files(self, limit: int, pattern: str = "**/*") -> Generator[FileAttachment, None, None]:
         """
-        Yields FileAttachments for files in the MemFS.
+        Yields FileAttachments for files found in the output directory.
 
         Args:
             limit: The maximum number of files to parse.
@@ -134,7 +144,7 @@ class MemFS:
         preload_dict: bool = False,
     ) -> list[FileAttachment]:
         """
-        Return a sorted list of output files found in the MemFS.
+        Return a sorted list of file paths within the output directory.
 
         Args:
             limit: The maximum number of files to parse.
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 4933a0a..c5f2f80 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -52,6 +52,8 @@ class NsJail:
         max_output_size: int = 1_000_000,
         read_chunk_size: int = 10_000,
         memfs_instance_size: int = 48 * Size.MiB,
+        memfs_home: str = "home",
+        memfs_output: str = "output",
         files_limit: int | None = 100,
         files_timeout: float | None = 8,
         files_pattern: str = "**/*",
@@ -65,6 +67,9 @@ class NsJail:
             max_output_size: Maximum size of the output in bytes.
             read_chunk_size: Size of the read buffer in bytes.
             memfs_instance_size: Size of the tmpfs instance in bytes.
+            memfs_home: Name of the mounted home directory.
+            memfs_output: Name of the output directory within home,
+                can be empty to use home as output.
             files_limit: Maximum number of output files to parse.
             files_timeout: Maximum time in seconds to wait for output files to be read.
             files_pattern: Pattern to match files to attach.
@@ -74,6 +79,8 @@ class NsJail:
         self.max_output_size = max_output_size
         self.read_chunk_size = read_chunk_size
         self.memfs_instance_size = memfs_instance_size
+        self.memfs_home = memfs_home
+        self.memfs_output = memfs_output
         self.files_limit = files_limit
         self.files_timeout = files_timeout
         self.files_pattern = files_pattern
@@ -191,7 +198,11 @@ class NsJail:
                 *nsjail_args,
             )
 
-        with NamedTemporaryFile() as nsj_log, MemFS(self.memfs_instance_size) as fs:
+        with NamedTemporaryFile() as nsj_log, MemFS(
+            instance_size=self.memfs_instance_size,
+            home=self.memfs_home,
+            output=self.memfs_output,
+        ) as fs:
             nsjail_args = (
                 # Mount `home` with Read/Write access
                 "--bindmount",