Add Filesystem abstraction.

fkdosilovic · fkdosilovic · commit 672df77c7356 · 2024-11-16T20:03:18.000+01:00
diff --git a/examples/0005_filesystem.py b/examples/0005_filesystem.py
@@ -0,0 +1,48 @@
+import judge0
+from judge0 import File, Filesystem
+
+print("Subexample 1")
+result = judge0.run(source_code="print('hello, world')")
+fs = Filesystem(result.post_execution_filesystem)
+for f in fs:
+    print(f.name)
+    print(f)
+    print()
+
+
+print("Subexample 2")
+fs = Filesystem(File("my_file.txt", "hello, world"))
+result = judge0.run(
+    source_code="print(open('my_file.txt').read())", additional_files=fs
+)
+print(result.stdout)
+for f in Filesystem(result.post_execution_filesystem):
+    print(f.name)
+    print(f)
+    print()
+
+
+print("Subexample 3")
+fs = Filesystem(File("my_file.txt", "hello, world"))
+result = judge0.run(
+    source_code="print(open('my_file.txt').read())", additional_files=fs
+)
+print(result.stdout)
+for f in result.post_execution_filesystem:
+    print(f.name)
+    print(f)
+    print()
+
+print("Subexample 4")
+fs = Filesystem(
+    [
+        File("my_file.txt", "hello, world"),
+        File("./dir1/dir2/dir3/my_file2.txt", "hello, world2"),
+    ]
+)
+result = judge0.run(source_code="find .", additional_files=fs, language=46)
+print(result.stdout)
+for f in Filesystem(result.post_execution_filesystem):
+    print(f.name)
+    print(f)
+    print()
diff --git a/src/judge0/__init__.py b/src/judge0/__init__.py
@@ -14,6 +14,7 @@
     SuluJudge0CE,
     SuluJudge0ExtraCE,
 )
+from .filesystem import File, Filesystem
 from .retry import MaxRetries, MaxWaitTime, RegularPeriodRetry
 from .submission import Submission
 
@@ -22,6 +23,8 @@
     "ATDJudge0CE",
     "ATDJudge0ExtraCE",
     "Client",
+    "File",
+    "Filesystem",
     "Language",
     "LanguageAlias",
     "MaxRetries",
diff --git a/src/judge0/base_types.py b/src/judge0/base_types.py
@@ -1,7 +1,14 @@
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import IntEnum
 
 
+class Encodeable(ABC):
+    @abstractmethod
+    def encode(self) -> bytes:
+        pass
+
+
 @dataclass
 class Language:
     id: int
diff --git a/src/judge0/common.py b/src/judge0/common.py
@@ -1,12 +1,16 @@
 from base64 import b64decode, b64encode
 from typing import Union
 
+from .base_types import Encodeable
 
-def encode(content: Union[bytes, str]) -> str:
+
+def encode(content: Union[bytes, str, Encodeable]) -> str:
     if isinstance(content, bytes):
         return b64encode(content).decode()
     if isinstance(content, str):
         return b64encode(content.encode()).decode()
+    if isinstance(content, Encodeable):
+        return b64encode(content.encode()).decode()
     raise ValueError(f"Unsupported type. Expected bytes or str, got {type(content)}!")
 
 
diff --git a/src/judge0/filesystem.py b/src/judge0/filesystem.py
@@ -0,0 +1,63 @@
+import copy
+import io
+import zipfile
+
+from base64 import b64decode, b64encode
+from collections import abc
+from typing import Iterable, Optional, Union
+
+from .base_types import Encodeable
+
+
+class File:
+    def __init__(self, name: str, content: Optional[Union[str, bytes]] = None):
+        self.name = name
+
+        # Let's keep content attribute internally encoded as bytes.
+        if isinstance(content, str):
+            self.content = content.encode()
+        elif isinstance(content, bytes):
+            self.content = content
+        else:
+            self.content = b""
+
+    def __str__(self):
+        return self.content.decode(errors="backslashreplace")
+
+
+class Filesystem(Encodeable):
+    def __init__(
+        self,
+        content: Optional[Union[str, bytes, File, Iterable[File], "Filesystem"]] = None,
+    ):
+        self.files: list[File] = []
+
+        if isinstance(content, (bytes, str)):
+            if isinstance(content, bytes):
+                zip_bytes = content
+            else:
+                zip_bytes = b64decode(content.encode())
+
+            with zipfile.ZipFile(io.BytesIO(zip_bytes), "r") as zip_file:
+                for file_name in zip_file.namelist():
+                    with zip_file.open(file_name) as fp:
+                        self.files.append(File(file_name, fp.read()))
+        elif isinstance(content, abc.Iterable):
+            self.files = list(content)
+        elif isinstance(content, File):
+            self.files = [content]
+        elif isinstance(content, Filesystem):
+            self.files = copy.deepcopy(content.files)
+
+    def encode(self) -> bytes:
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
+            for file in self.files:
+                zip_file.writestr(file.name, file.content)
+        return zip_buffer.getvalue()
+
+    def __str__(self) -> str:
+        return b64encode(self.encode()).decode()
+
+    def __iter__(self):
+        return iter(self.files)
diff --git a/src/judge0/submission.py b/src/judge0/submission.py
@@ -1,6 +1,8 @@
 from datetime import datetime
 from typing import Union
 
+from judge0.filesystem import Filesystem
+
 from .base_types import LanguageAlias, Status
 from .common import decode, encode
 
@@ -10,7 +12,12 @@
     "stdin",
     "expected_output",
 }
-ENCODED_RESPONSE_FIELDS = {"stdout", "stderr", "compile_output"}
+ENCODED_RESPONSE_FIELDS = {
+    "stdout",
+    "stderr",
+    "compile_output",
+    # "post_execution_filesystem",
+}
 ENCODED_FIELDS = ENCODED_REQUEST_FIELDS | ENCODED_RESPONSE_FIELDS
 EXTRA_REQUEST_FIELDS = {
     "compiler_options",
@@ -141,6 +148,8 @@ def set_attributes(self, attributes):
                 value = datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%fZ")
             elif attr in FLOATING_POINT_FIELDS and value is not None:
                 value = float(value)
+            elif attr == "post_execution_filesystem":
+                value = Filesystem(value)
 
             setattr(self, attr, value)
 
