Extract archives to a single zip file

Author: PeterJCLaw

code_submitter/extract_archives.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import asyncio
+import zipfile
 import argparse
 from pathlib import Path
 
@@ -9,33 +10,24 @@
 from . import utils, config
 
 
-async def async_main(submissions_directory: Path) -> None:
-    submissions_directory.mkdir(parents=True, exist_ok=True)
+async def async_main(output_archive: Path) -> None:
+    output_archive.parent.mkdir(parents=True, exist_ok=True)
 
     database = databases.Database(config.DATABASE_URL)
 
-    async with database.transaction():
-        submissions = await utils.get_chosen_submissions(database)
-
-        for team, (_, content) in submissions.items():
-            (submissions_directory / f'{team.upper()}.zip').write_bytes(content)
-
-        (submissions_directory / 'summary.txt').write_text(
-            "".join(
-                "{}: {}\n".format(team, id_)
-                for team, (id_, _) in submissions.items()
-            ),
-        )
+    with zipfile.ZipFile(output_archive) as zf:
+        async with database.transaction():
+            utils.collect_submissions(database, zf)
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
-    parser.add_argument('submissions_directory', type=Path)
+    parser.add_argument('output_archive', type=Path)
     return parser.parse_args()
 
 
 def main(args: argparse.Namespace) -> None:
-    asyncio.get_event_loop().run_until_complete(async_main(args.submissions_directory))
+    asyncio.get_event_loop().run_until_complete(async_main(args.output_archive))
 
 
 if __name__ == '__main__':

code_submitter/utils.py

@@ -1,4 +1,5 @@
 from typing import Dict, Tuple
+from zipfile import ZipFile
 
 import databases
 from sqlalchemy.sql import select
@@ -32,3 +33,22 @@ async def get_chosen_submissions(
 
     # Rely on later keys replacing earlier occurrences of the same key.
     return {x['team']: (x['id'], x['content']) for x in rows}
+
+
+def summarise(submissions: Dict[str, Tuple[int, bytes]]) -> str:
+    return "".join(
+        "{}: {}\n".format(team, id_)
+        for team, (id_, _) in sorted(submissions.items())
+    )
+
+
+async def collect_submissions(
+    database: databases.Database,
+    zipfile: ZipFile,
+) -> None:
+    submissions = await get_chosen_submissions(database)
+
+    for team, (_, content) in submissions.items():
+        zipfile.writestr(f'{team.upper()}.zip', content)
+
+    zipfile.writestr('summary.txt', summarise(submissions))

tests/tests_utils.py

@@ -1,3 +1,5 @@
+import io
+import zipfile
 import datetime
 
 import test_utils
@@ -72,3 +74,38 @@ def test_get_chosen_submissions_multiple_chosen(self) -> None:
             },
             result,
         )
+
+    def test_collect_submissions(self) -> None:
+        self.await_(self.database.execute(
+            ChoiceHistory.insert().values(
+                archive_id=8888888888,
+                username='someone_else',
+                created=datetime.datetime(2020, 8, 8, 12, 0),
+            ),
+        ))
+        self.await_(self.database.execute(
+            ChoiceHistory.insert().values(
+                archive_id=1111111111,
+                username='test_user',
+                created=datetime.datetime(2020, 3, 3, 12, 0),
+            ),
+        ))
+        self.await_(self.database.execute(
+            ChoiceHistory.insert().values(
+                archive_id=2222222222,
+                username='test_user',
+                created=datetime.datetime(2020, 2, 2, 12, 0),
+            ),
+        ))
+
+        with zipfile.ZipFile(io.BytesIO(), mode='w') as zf:
+            self.await_(utils.collect_submissions(self.database, zf))
+
+            self.assertEqual(
+                {
+                    'summary.txt': b'ABC: 8888888888\nSRZ2: 1111111111\n',
+                    'SRZ2.zip': b'1111111111',
+                    'ABC.zip': b'8888888888',
+                },
+                {x: zf.open(x).read() for x in zf.namelist()},
+            )