diff --git a/gallery_dl/postprocessor/classify.py b/gallery_dl/postprocessor/classify.py index 34af1d94..56429559 100644 --- a/gallery_dl/postprocessor/classify.py +++ b/gallery_dl/postprocessor/classify.py @@ -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"), - "Video" : ("flv", "ogv", "avi", "mp4", "mpg", "mpeg", "3gp", "mkv", - "webm", "vob", "wmv"), - "Pictures" : ("jpg", "jpeg", "png", "gif", "bmp", "svg", "webp"), + "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", "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 diff --git a/test/test_postprocessor.py b/test/test_postprocessor.py index dd53803f..2941b815 100644 --- a/test/test_postprocessor.py +++ b/test/test_postprocessor.py @@ -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):