Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
mikf
GitHub Repository: mikf/gallery-dl
Path: blob/master/gallery_dl/postprocessor/compare.py
8753 views
1
# -*- coding: utf-8 -*-
2
3
# Copyright 2020-2025 Mike Fährmann
4
#
5
# This program is free software; you can redistribute it and/or modify
6
# it under the terms of the GNU General Public License version 2 as
7
# published by the Free Software Foundation.
8
9
"""Compare versions of the same file and replace/enumerate them on mismatch"""
10
11
from .common import PostProcessor
12
from .. import text, util, output, exception
13
import os
14
15
16
class ComparePP(PostProcessor):
17
18
def __init__(self, job, options):
19
PostProcessor.__init__(self, job)
20
if options.get("shallow"):
21
self._compare = self._compare_size
22
self._equal_exc = self._equal_cnt = 0
23
24
if equal := options.get("equal"):
25
equal, _, emax = equal.partition(":")
26
self._equal_max = text.parse_int(emax)
27
if equal == "abort":
28
self._equal_exc = exception.StopExtraction
29
elif equal == "terminate":
30
self._equal_exc = exception.TerminateExtraction
31
elif equal == "exit":
32
self._equal_exc = SystemExit
33
34
job.register_hooks({"file": (
35
self.enumerate
36
if options.get("action") == "enumerate" else
37
self.replace
38
)}, options)
39
40
def replace(self, pathfmt):
41
try:
42
if self._compare(pathfmt.realpath, pathfmt.temppath):
43
return self._equal(pathfmt)
44
except OSError:
45
pass
46
self._equal_cnt = 0
47
48
def enumerate(self, pathfmt):
49
num = 1
50
try:
51
while not self._compare(pathfmt.realpath, pathfmt.temppath):
52
pathfmt.prefix = prefix = format(num) + "."
53
pathfmt.kwdict["extension"] = prefix + pathfmt.extension
54
pathfmt.build_path()
55
num += 1
56
return self._equal(pathfmt)
57
except OSError:
58
pass
59
self._equal_cnt = 0
60
61
def _compare(self, f1, f2):
62
return self._compare_size(f1, f2) and self._compare_content(f1, f2)
63
64
def _compare_size(self, f1, f2):
65
return os.stat(f1).st_size == os.stat(f2).st_size
66
67
def _compare_content(self, f1, f2):
68
size = 16384
69
with open(f1, "rb") as fp1, open(f2, "rb") as fp2:
70
while True:
71
buf1 = fp1.read(size)
72
buf2 = fp2.read(size)
73
if buf1 != buf2:
74
return False
75
if not buf1:
76
return True
77
78
def _equal(self, pathfmt):
79
if self._equal_exc:
80
self._equal_cnt += 1
81
if self._equal_cnt >= self._equal_max:
82
util.remove_file(pathfmt.temppath)
83
output.stderr_write("\n")
84
raise self._equal_exc()
85
pathfmt.delete = True
86
87
88
__postprocessor__ = ComparePP
89
90