1
0
mirror of https://github.com/mikf/gallery-dl.git synced 2024-11-21 18:22:30 +01:00

[pp:classify] rewrite & simplify (#5213)

Do not manually build paths, which get later overwritten by
pathfmt.build_path() anyway. Just update the target directory and let
the rest of the "path logic" handle it.

Fixes skipping previously downloaded and categorized files,
which was broken since 8124c16a50
This commit is contained in:
Mike Fährmann 2024-11-18 20:08:15 +01:00
parent 00fe1c81b2
commit 9b2d782cb7
No known key found for this signature in database
GPG Key ID: 5680CA389D365A88
2 changed files with 48 additions and 42 deletions

View File

@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
# Copyright 2018-2021 Mike Fährmann
# Copyright 2018-2024 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
@ -15,44 +15,43 @@ import os
class ClassifyPP(PostProcessor):
DEFAULT_MAPPING = {
"Music" : ("mp3", "aac", "flac", "ogg", "wma", "m4a", "wav"),
"Pictures" : ("jpg", "jpeg", "png", "gif", "bmp", "svg", "webp",
"avif", "heic", "heif", "ico", "psd"),
"Video" : ("flv", "ogv", "avi", "mp4", "mpg", "mpeg", "3gp", "mkv",
"webm", "vob", "wmv"),
"Pictures" : ("jpg", "jpeg", "png", "gif", "bmp", "svg", "webp"),
"webm", "vob", "wmv", "m4v", "mov"),
"Music" : ("mp3", "aac", "flac", "ogg", "wma", "m4a", "wav"),
"Archives" : ("zip", "rar", "7z", "tar", "gz", "bz2"),
"Documents": ("txt", "pdf"),
}
def __init__(self, job, options):
PostProcessor.__init__(self, job)
mapping = options.get("mapping", self.DEFAULT_MAPPING)
self.directory = self.realdirectory = ""
mapping = options.get("mapping", self.DEFAULT_MAPPING)
self.mapping = {
ext: directory
for directory, exts in mapping.items()
for ext in exts
}
job.register_hooks(
{"prepare": self.prepare, "file": self.move}, options)
job.register_hooks({
"post" : self.initialize,
"prepare": self.prepare,
}, options)
def initialize(self, pathfmt):
# store base directory paths
self.directory = pathfmt.directory
self.realdirectory = pathfmt.realdirectory
def prepare(self, pathfmt):
# extend directory paths depending on file extension
ext = pathfmt.extension
if ext in self.mapping:
# set initial paths to enable download skips
self._build_paths(pathfmt, self.mapping[ext])
def move(self, pathfmt):
ext = pathfmt.extension
if ext in self.mapping:
# rebuild paths in case the filename extension changed
path = self._build_paths(pathfmt, self.mapping[ext])
os.makedirs(path, exist_ok=True)
@staticmethod
def _build_paths(pathfmt, extra):
path = pathfmt.realdirectory + extra
pathfmt.realpath = path + os.sep + pathfmt.filename
pathfmt.path = pathfmt.directory + extra + os.sep + pathfmt.filename
return path
extra = self.mapping[ext] + os.sep
pathfmt.directory = self.directory + extra
pathfmt.realdirectory = self.realdirectory + extra
__postprocessor__ = ClassifyPP

View File

@ -120,30 +120,37 @@ class ClassifyTest(BasePostprocessorTest):
for directory, exts in pp.DEFAULT_MAPPING.items()
for ext in exts
})
self.pathfmt.set_extension("jpg")
self.pathfmt.build_path()
pp.prepare(self.pathfmt)
self.assertEqual(pp.directory, "")
self._trigger(("post",))
self.assertEqual(pp.directory, self.pathfmt.directory)
self.pathfmt.set_extension("jpg")
self._trigger(("prepare",))
self.pathfmt.build_path()
path = os.path.join(self.dir.name, "test", "Pictures")
self.assertEqual(self.pathfmt.path, path + "/file.jpg")
self.assertEqual(self.pathfmt.realpath, path + "/file.jpg")
with patch("os.makedirs") as mkdirs:
self._trigger()
mkdirs.assert_called_once_with(path, exist_ok=True)
self.pathfmt.set_extension("mp4")
self._trigger(("prepare",))
self.pathfmt.build_path()
path = os.path.join(self.dir.name, "test", "Video")
self.assertEqual(self.pathfmt.path, path + "/file.mp4")
self.assertEqual(self.pathfmt.realpath, path + "/file.mp4")
def test_classify_noop(self):
pp = self._create()
rp = self.pathfmt.realpath
pp.prepare(self.pathfmt)
self.assertEqual(pp.directory, "")
self._trigger(("post",))
self._trigger(("prepare",))
self.assertEqual(pp.directory, self.pathfmt.directory)
self.assertEqual(self.pathfmt.path, rp)
self.assertEqual(self.pathfmt.realpath, rp)
with patch("os.makedirs") as mkdirs:
self._trigger()
self.assertEqual(mkdirs.call_count, 0)
def test_classify_custom(self):
pp = self._create({"mapping": {
"foo/bar": ["foo", "bar"],
@ -153,18 +160,18 @@ class ClassifyTest(BasePostprocessorTest):
"foo": "foo/bar",
"bar": "foo/bar",
})
self.pathfmt.set_extension("foo")
self.pathfmt.build_path()
pp.prepare(self.pathfmt)
self.assertEqual(pp.directory, "")
self._trigger(("post",))
self.assertEqual(pp.directory, self.pathfmt.directory)
self.pathfmt.set_extension("foo")
self._trigger(("prepare",))
self.pathfmt.build_path()
path = os.path.join(self.dir.name, "test", "foo", "bar")
self.assertEqual(self.pathfmt.path, path + "/file.foo")
self.assertEqual(self.pathfmt.realpath, path + "/file.foo")
with patch("os.makedirs") as mkdirs:
self._trigger()
mkdirs.assert_called_once_with(path, exist_ok=True)
class ExecTest(BasePostprocessorTest):