1
0
mirror of https://github.com/mikf/gallery-dl.git synced 2024-11-22 18:53:21 +01:00
gallery-dl/gallery_dl/job.py
Mike Fährmann 48a8717a7c
add 'output.num-to-str' option
... to convert any numeric values to string when outputting them as JSON
(during '--dump-json' or otherwise)
2018-10-08 20:28:54 +02:00

538 lines
17 KiB
Python

# -*- coding: utf-8 -*-
# Copyright 2015-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.
import sys
import time
import json
import hashlib
import logging
from . import extractor, downloader, postprocessor
from . import config, util, output, exception
from .extractor.message import Message
class Job():
"""Base class for Job-types"""
ulog = None
def __init__(self, url, parent=None):
self.url = url
self.extractor = extractor.find(url)
if self.extractor is None:
raise exception.NoExtractorError(url)
self.extractor.log.debug(
"Using %s for '%s'", self.extractor.__class__.__name__, url)
# url predicates
self.pred_url = self._prepare_predicates(
"image", [util.UniquePredicate()], True)
# queue predicates
self.pred_queue = self._prepare_predicates(
"chapter", [], False)
# category transfer
if parent and parent.extractor.categorytransfer:
self.extractor.category = parent.extractor.category
self.extractor.subcategory = parent.extractor.subcategory
# user-supplied metadata
self.userkwds = self.extractor.config("keywords")
def run(self):
"""Execute or run the job"""
try:
log = self.extractor.log
for msg in self.extractor:
self.dispatch(msg)
except exception.AuthenticationError as exc:
msg = str(exc) or "Please provide a valid username/password pair."
log.error("Authentication failed: %s", msg)
except exception.AuthorizationError:
log.error("You do not have permission to access the resource "
"at '%s'", self.url)
except exception.NotFoundError as exc:
res = str(exc) or "resource (gallery/image/user)"
log.error("The %s at '%s' does not exist", res, self.url)
except exception.HttpError as exc:
log.error("HTTP request failed: %s", exc)
except exception.FormatError as exc:
err, obj = exc.args
log.error("Applying %s format string failed: %s: %s",
obj, err.__class__.__name__, err)
except exception.FilterError as exc:
err = exc.args[0]
log.error("Evaluating filter expression failed: %s: %s",
err.__class__.__name__, err)
except exception.StopExtraction:
pass
except OSError as exc:
log.error("Unable to download data: %s", exc)
except Exception as exc:
log.error(("An unexpected error occurred: %s - %s. "
"Please run gallery-dl again with the --verbose flag, "
"copy its output and report this issue on "
"https://github.com/mikf/gallery-dl/issues ."),
exc.__class__.__name__, exc)
log.debug("Traceback", exc_info=True)
self.handle_finalize()
def dispatch(self, msg):
"""Call the appropriate message handler"""
if msg[0] == Message.Url:
_, url, kwds = msg
if self.pred_url(url, kwds):
self.update_kwdict(kwds)
self.handle_url(url, kwds)
elif msg[0] == Message.Directory:
self.update_kwdict(msg[1])
self.handle_directory(msg[1])
elif msg[0] == Message.Queue:
_, url, kwds = msg
if self.pred_queue(url, kwds):
self.handle_queue(url, kwds)
elif msg[0] == Message.Urllist:
_, urls, kwds = msg
if self.pred_url(urls[0], kwds):
self.update_kwdict(kwds)
self.handle_urllist(urls, kwds)
elif msg[0] == Message.Version:
if msg[1] != 1:
raise "unsupported message-version ({}, {})".format(
self.extractor.category, msg[1]
)
# TODO: support for multiple message versions
def handle_url(self, url, keywords):
"""Handle Message.Url"""
def handle_urllist(self, urls, keywords):
"""Handle Message.Urllist"""
self.handle_url(urls[0], keywords)
def handle_directory(self, keywords):
"""Handle Message.Directory"""
def handle_queue(self, url, keywords):
"""Handle Message.Queue"""
def handle_finalize(self):
"""Handle job finalization"""
def update_kwdict(self, kwdict):
"""Update 'kwdict' with additional metadata"""
kwdict["category"] = self.extractor.category
kwdict["subcategory"] = self.extractor.subcategory
if self.userkwds:
kwdict.update(self.userkwds)
def _prepare_predicates(self, target, predicates, skip=True):
pfilter = self.extractor.config(target + "-filter")
if pfilter:
try:
pred = util.FilterPredicate(pfilter, target)
except (SyntaxError, ValueError, TypeError) as exc:
self.extractor.log.warning(exc)
else:
predicates.append(pred)
prange = self.extractor.config(target + "-range")
if prange:
try:
pred = util.RangePredicate(prange)
except ValueError as exc:
self.extractor.log.warning(
"invalid %s range: %s", target, exc)
else:
if skip and pred.lower > 1 and not pfilter:
pred.index += self.extractor.skip(pred.lower - 1)
predicates.append(pred)
return util.build_predicate(predicates)
def _write_unsupported(self, url):
if self.ulog:
self.ulog.info(url)
class DownloadJob(Job):
"""Download images into appropriate directory/filename locations"""
def __init__(self, url, parent=None):
Job.__init__(self, url, parent)
self.log = logging.getLogger("download")
self.pathfmt = None
self.archive = None
self.sleep = None
self.downloaders = {}
self.postprocessors = None
self.out = output.select()
def handle_url(self, url, keywords, fallback=None):
"""Download the resource specified in 'url'"""
# prepare download
self.pathfmt.set_keywords(keywords)
if self.pathfmt.exists(self.archive):
self.out.skip(self.pathfmt.path)
return
if self.sleep:
time.sleep(self.sleep)
# download from URL
if not self.download(url):
# use fallback URLs if available
for num, url in enumerate(fallback or (), 1):
self.log.info("Trying fallback URL #%d", num)
if self.download(url):
break
else:
# download failed
self.log.error(
"Failed to download %s", self.pathfmt.filename or url)
return
if not self.pathfmt.temppath:
self.out.skip(self.pathfmt.path)
return
# run post processors
if self.postprocessors:
for pp in self.postprocessors:
pp.run(self.pathfmt)
# download succeeded
self.pathfmt.finalize()
self.out.success(self.pathfmt.path, 0)
if self.archive:
self.archive.add(keywords)
def handle_urllist(self, urls, keywords):
"""Download the resource specified in 'url'"""
fallback = iter(urls)
url = next(fallback)
self.handle_url(url, keywords, fallback)
def handle_directory(self, keywords):
"""Set and create the target directory for downloads"""
if not self.pathfmt:
self.initialize()
self.pathfmt.set_directory(keywords)
def handle_queue(self, url, keywords):
try:
self.__class__(url, self).run()
except exception.NoExtractorError:
self._write_unsupported(url)
def handle_finalize(self):
if self.postprocessors:
for pp in self.postprocessors:
pp.finalize()
def download(self, url):
"""Download 'url'"""
scheme = url.partition(":")[0]
downloader = self.get_downloader(scheme)
if downloader:
return downloader.download(url, self.pathfmt)
return False
def get_downloader(self, scheme):
"""Return a downloader suitable for 'scheme'"""
if scheme == "https":
scheme = "http"
try:
return self.downloaders[scheme]
except KeyError:
pass
klass = downloader.find(scheme)
if klass:
instance = klass(self.extractor, self.out)
else:
instance = None
self.log.error("'%s:' URLs are not supported", scheme)
self.downloaders[scheme] = instance
return instance
def initialize(self):
"""Delayed initialization of PathFormat, etc."""
self.pathfmt = util.PathFormat(self.extractor)
self.sleep = self.extractor.config("sleep")
archive = self.extractor.config("archive")
if archive:
path = util.expand_path(archive)
self.archive = util.DownloadArchive(path, self.extractor)
postprocessors = self.extractor.config("postprocessors")
if postprocessors:
self.postprocessors = []
for pp_dict in postprocessors:
whitelist = pp_dict.get("whitelist")
blacklist = pp_dict.get("blacklist")
if (whitelist and self.extractor.category not in whitelist or
blacklist and self.extractor.category in blacklist):
continue
name = pp_dict.get("name")
pp_cls = postprocessor.find(name)
if not pp_cls:
postprocessor.log.warning("module '%s' not found", name)
continue
try:
pp_obj = pp_cls(self.pathfmt, pp_dict)
except Exception as exc:
postprocessor.log.error(
"%s: initialization failed: %s %s",
name, exc.__class__.__name__, exc)
else:
self.postprocessors.append(pp_obj)
self.extractor.log.debug(
"Active postprocessor modules: %s", self.postprocessors)
class SimulationJob(DownloadJob):
"""Simulate the extraction process without downloading anything"""
def handle_url(self, url, keywords, fallback=None):
self.pathfmt.set_keywords(keywords)
self.out.skip(self.pathfmt.path)
if self.sleep:
time.sleep(self.sleep)
if self.archive:
self.archive.add(keywords)
def handle_directory(self, keywords):
if not self.pathfmt:
self.initialize()
class KeywordJob(Job):
"""Print available keywords"""
def handle_url(self, url, keywords):
print("\nKeywords for filenames and --filter:")
print("------------------------------------")
self.print_keywords(keywords)
raise exception.StopExtraction()
def handle_directory(self, keywords):
print("Keywords for directory names:")
print("-----------------------------")
self.print_keywords(keywords)
def handle_queue(self, url, keywords):
if not keywords:
self.extractor.log.info(
"This extractor delegates work to other extractors "
"and does not provide any keywords on its own. Try "
"'gallery-dl -K \"%s\"' instead.", url)
else:
print("Keywords for --chapter-filter:")
print("------------------------------")
self.print_keywords(keywords)
if self.extractor.categorytransfer:
print()
KeywordJob(url, self).run()
raise exception.StopExtraction()
@staticmethod
def print_keywords(keywords, prefix=""):
"""Print key-value pairs with formatting"""
suffix = "]" if prefix else ""
for key, value in sorted(keywords.items()):
key = prefix + key + suffix
if isinstance(value, dict):
KeywordJob.print_keywords(value, key + "[")
elif isinstance(value, list):
if value and isinstance(value[0], dict):
KeywordJob.print_keywords(value[0], key + "[][")
else:
print(key, "[]", sep="")
for val in value:
print(" -", val)
else:
# string or number
print(key, "\n ", value, sep="")
class UrlJob(Job):
"""Print download urls"""
maxdepth = 1
def __init__(self, url, parent=None, depth=1):
Job.__init__(self, url, parent)
self.depth = depth
if depth >= self.maxdepth:
self.handle_queue = self.handle_url
@staticmethod
def handle_url(url, _):
print(url)
@staticmethod
def handle_urllist(urls, _):
prefix = ""
for url in urls:
print(prefix, url, sep="")
prefix = "| "
def handle_queue(self, url, _):
try:
UrlJob(url, self, self.depth + 1).run()
except exception.NoExtractorError:
self._write_unsupported(url)
class TestJob(DownloadJob):
"""Generate test-results for extractor runs"""
class HashIO():
"""Minimal file-like interface"""
def __init__(self, hashobj):
self.hashobj = hashobj
self.path = ""
self.size = 0
self.has_extension = True
def __enter__(self):
return self
def __exit__(self, *args):
pass
def open(self, mode):
self.size = 0
return self
def write(self, content):
"""Update SHA1 hash"""
self.size += len(content)
self.hashobj.update(content)
def tell(self):
return self.size
def part_size(self):
return 0
def __init__(self, url, parent=None, content=False):
DownloadJob.__init__(self, url, parent)
self.content = content
self.list_url = []
self.list_keyword = []
self.list_archive = []
self.hash_url = hashlib.sha1()
self.hash_keyword = hashlib.sha1()
self.hash_archive = hashlib.sha1()
self.hash_content = hashlib.sha1()
if content:
self.fileobj = self.HashIO(self.hash_content)
self.get_downloader("http")._check_extension = lambda a, b: None
def run(self):
for msg in self.extractor:
self.dispatch(msg)
def handle_url(self, url, keywords):
self.update_url(url)
self.update_keyword(keywords)
self.update_archive(keywords)
self.update_content(url)
def handle_urllist(self, urls, keywords):
self.handle_url(urls[0], keywords)
def handle_directory(self, keywords):
self.update_keyword(keywords, False)
def handle_queue(self, url, keywords):
self.update_url(url)
self.update_keyword(keywords)
def update_url(self, url):
"""Update the URL hash"""
self.list_url.append(url)
self.hash_url.update(url.encode())
def update_keyword(self, kwdict, to_list=True):
"""Update the keyword hash"""
if to_list:
self.list_keyword.append(kwdict.copy())
self.hash_keyword.update(
json.dumps(kwdict, sort_keys=True).encode())
def update_archive(self, kwdict):
"""Update the archive-id hash"""
archive_id = self.extractor.archive_fmt.format_map(kwdict)
self.list_archive.append(archive_id)
self.hash_archive.update(archive_id.encode())
def update_content(self, url):
"""Update the content hash"""
if self.content:
scheme = url.partition(":")[0]
self.get_downloader(scheme).download(url, self.fileobj)
class DataJob(Job):
"""Collect extractor results and dump them"""
def __init__(self, url, parent=None, file=sys.stdout):
Job.__init__(self, url, parent)
self.file = file
self.data = []
def run(self):
# collect data
try:
for msg in self.extractor:
self.dispatch(msg)
except Exception as exc:
self.data.append((exc.__class__.__name__, str(exc)))
except BaseException:
pass
if config.get(("output", "num-to-str"), False):
for msg in self.data:
util.transform_dict(msg[-1], util.number_to_string)
# dump to 'file'
json.dump(
self.data, self.file,
sort_keys=True, indent=2,
ensure_ascii=config.get(("output", "ascii"), True),
)
self.file.write("\n")
def handle_url(self, url, keywords):
self.data.append((Message.Url, url, keywords.copy()))
def handle_urllist(self, urls, keywords):
self.data.append((Message.Urllist, list(urls), keywords.copy()))
def handle_directory(self, keywords):
self.data.append((Message.Directory, keywords.copy()))
def handle_queue(self, url, keywords):
self.data.append((Message.Queue, url, keywords.copy()))
def handle_finalize(self):
self.file.close()