diff --git a/README.md b/README.md index 12a47643..08de5602 100644 --- a/README.md +++ b/README.md @@ -347,8 +347,5 @@ submission = Submission( result = judge0.run(submissions=submission) print(result.stdout) - -matches = [f for f in result.post_execution_filesystem if f.name == "my_dir2/my_file2.txt"] -f = matches[0] if matches else None -print(f) +print(result.post_execution_filesystem.find("./my_dir2/my_file2.txt")) ``` diff --git a/src/judge0/filesystem.py b/src/judge0/filesystem.py index 27fae223..b753dfc4 100644 --- a/src/judge0/filesystem.py +++ b/src/judge0/filesystem.py @@ -92,6 +92,27 @@ def encode(self) -> bytes: zip_file.writestr(file.name, file.content) return zip_buffer.getvalue() + def find(self, name: str) -> Optional[File]: + """Find file by name in Filesystem object. + + Parameters + ---------- + name : str + File name to find. + + Returns + ------- + File or None + Found File object or None if not found. + """ + if name.startswith("./"): + name = name[2:] + elif name.startswith("/"): + name = name[1:] + + matches = [f for f in self.files if f.name == name] + return matches[0] if matches else None + def __str__(self) -> str: """Create string representation of Filesystem object.""" return b64encode(self.encode()).decode()