From 3313a936debe04c733b3003858724de17cfcd7de Mon Sep 17 00:00:00 2001
From: jwansek <eddie.atten.ea29@gmail.com>
Date: Sat, 22 Jan 2022 20:42:03 +0000
Subject: added automatically removing generated files

---
 mark.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

(limited to 'mark.py')

diff --git a/mark.py b/mark.py
index 2630be1..835ac03 100644
--- a/mark.py
+++ b/mark.py
@@ -27,15 +27,22 @@ class FileDependencies:
             pass
 
     def __exit__(self, type, value, traceback):
+        stuff_to_remove = []
         try:
-            for file_dep in self.assessment_struct["dependencies"]["files"]:
-                if os.path.isfile(os.path.split(file_dep)[-1]):
-                    os.remove(os.path.split(file_dep)[-1])
-                else:
-                    shutil.rmtree(os.path.split(file_dep)[-1])
+            stuff_to_remove += [os.path.split(f)[-1] for f in self.assessment_struct["dependencies"]["files"]]
+        except KeyError:
+            pass
+        try:
+            stuff_to_remove += self.assessment_struct["produced_files"]
         except KeyError:
             pass
 
+        for file_dep in stuff_to_remove:
+            if os.path.isfile(file_dep):
+                os.remove(file_dep)
+            else:
+                shutil.rmtree(file_dep)
+
 def main(**kwargs):
     student_no = os.path.splitext(os.path.split(args["submission"])[-1])[0]
 
-- 
cgit v1.2.3