From b36125333f0d465af48bd1a2d8b9c83a2e08bf84 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20F=C3=A4hrmann?= Date: Fri, 9 Sep 2022 11:41:27 +0200 Subject: [PATCH] [postprocessor:zip] implement 'files' option (#2872) --- docs/configuration.rst | 13 +++++++++++++ gallery_dl/postprocessor/zip.py | 19 +++++++++++++++++++ test/test_postprocessor.py | 16 ++++++++++------ 3 files changed, 42 insertions(+), 6 deletions(-) diff --git a/docs/configuration.rst b/docs/configuration.rst index 7ab49275..44853e3b 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -3805,6 +3805,19 @@ Description Filename extension for the created ZIP archive. +zip.files +--------- +Type + ``list`` of |Path| +Example + ``["info.json"]`` +Description + List of extra files to be added to a ZIP archive. + + Note: Relative paths are relative to the current + `download directory `__. + + zip.keep-files -------------- Type diff --git a/gallery_dl/postprocessor/zip.py b/gallery_dl/postprocessor/zip.py index ff97add2..4f376fe6 100644 --- a/gallery_dl/postprocessor/zip.py +++ b/gallery_dl/postprocessor/zip.py @@ -26,6 +26,7 @@ class ZipPP(PostProcessor): def __init__(self, job, options): PostProcessor.__init__(self, job) self.delete = not options.get("keep-files", False) + self.files = options.get("files") ext = "." + options.get("extension", "zip") algorithm = options.get("compression", "store") if algorithm not in self.COMPRESSION_ALGORITHMS: @@ -56,6 +57,9 @@ class ZipPP(PostProcessor): # 'NameToInfo' is not officially documented, but it's available # for all supported Python versions and using it directly is a lot # faster than calling getinfo() + if self.files: + self.write_extra(pathfmt, zfile, self.files) + self.files = None if pathfmt.filename not in zfile.NameToInfo: zfile.write(pathfmt.temppath, pathfmt.filename) pathfmt.delete = self.delete @@ -69,6 +73,21 @@ class ZipPP(PostProcessor): with self.open() as zfile: self.write(pathfmt, zfile) + def write_extra(self, pathfmt, zfile, files): + for path in map(util.expand_path, files): + if not os.path.isabs(path): + path = os.path.join(pathfmt.realdirectory, path) + try: + zfile.write(path, os.path.basename(path)) + except OSError as exc: + self.log.warning( + "Unable to write %s to %s", path, zfile.filename) + self.log.debug("%s: %s", exc, exc.__class__.__name__) + pass + else: + if self.delete: + util.remove_file(path) + def finalize(self, pathfmt, status): if self.zfile: self.zfile.close() diff --git a/test/test_postprocessor.py b/test/test_postprocessor.py index 42babd30..af8b0afc 100644 --- a/test/test_postprocessor.py +++ b/test/test_postprocessor.py @@ -452,9 +452,11 @@ class ZipTest(BasePostprocessorTest): self.assertTrue(pp.args[0].endswith("/test.cbz")) def test_zip_write(self): - pp = self._create() - with tempfile.NamedTemporaryFile("w", dir=self.dir.name) as file: + pp = self._create({"files": [file.name, "_info_.json"], + "keep-files": True}) + + filename = os.path.basename(file.name) file.write("foobar\n") # write dummy file with 3 different names @@ -466,18 +468,19 @@ class ZipTest(BasePostprocessorTest): self._trigger() nti = pp.zfile.NameToInfo - self.assertEqual(len(nti), i+1) + self.assertEqual(len(nti), i+2) self.assertIn(name, nti) # check file contents - self.assertEqual(len(nti), 3) + self.assertEqual(len(nti), 4) self.assertIn("file0.ext", nti) self.assertIn("file1.ext", nti) self.assertIn("file2.ext", nti) + self.assertIn(filename, nti) # write the last file a second time (will be skipped) self._trigger() - self.assertEqual(len(pp.zfile.NameToInfo), 3) + self.assertEqual(len(pp.zfile.NameToInfo), 4) # close file self._trigger(("finalize",), 0) @@ -485,10 +488,11 @@ class ZipTest(BasePostprocessorTest): # reopen to check persistence with zipfile.ZipFile(pp.zfile.filename) as file: nti = file.NameToInfo - self.assertEqual(len(pp.zfile.NameToInfo), 3) + self.assertEqual(len(pp.zfile.NameToInfo), 4) self.assertIn("file0.ext", nti) self.assertIn("file1.ext", nti) self.assertIn("file2.ext", nti) + self.assertIn(filename, nti) os.unlink(pp.zfile.filename)