summaryrefslogtreecommitdiffstats
path: root/gallery_dl/postprocessor/zip.py
blob: 3a0c3237188533053052dfb6041bb610d7d47734 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# -*- coding: utf-8 -*-

# Copyright 2018 Mike Fährmann
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.

"""Store files in ZIP archives"""

from .common import PostProcessor
import zipfile
import os


class ZipPP(PostProcessor):

    COMPRESSION_ALGORITHMS = {
        "store": zipfile.ZIP_STORED,
        "zip": zipfile.ZIP_DEFLATED,
        "bzip2": zipfile.ZIP_BZIP2,
        "lzma": zipfile.ZIP_LZMA,
    }

    def __init__(self, pathfmt, options):
        PostProcessor.__init__(self)
        self.delete = not options.get("keep-files", False)
        self.ext = "." + options.get("extension", "zip")
        algorithm = options.get("compression", "store")
        if algorithm not in self.COMPRESSION_ALGORITHMS:
            self.log.warning(
                "unknown compression algorithm '%s'; falling back to 'store'",
                algorithm)
            algorithm = "store"

        self.path = pathfmt.realdirectory
        self.zfile = zipfile.ZipFile(
            self.path + self.ext, "a",
            self.COMPRESSION_ALGORITHMS[algorithm], True)

    def run(self, pathfmt):
        # 'NameToInfo' is not officially documented, but it's available
        # for all supported Python versions and using it directly is a lot
        # better than calling getinfo()
        if pathfmt.filename not in self.zfile.NameToInfo:
            self.zfile.write(pathfmt.temppath, pathfmt.filename)
            pathfmt.delete = self.delete

    def finalize(self):
        self.zfile.close()

        if self.delete:
            try:
                os.rmdir(self.path)
            except OSError:
                pass

            if not self.zfile.NameToInfo:
                try:
                    os.unlink(self.zfile.filename)
                except OSError:
                    pass


__postprocessor__ = ZipPP