From e9ad07791cff0f12e280d2959b58cbc23c2cdd95 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Fri, 7 Nov 2014 21:17:28 +0530 Subject: [PATCH] Start work on new save to disk implementation --- src/calibre/gui2/save.py | 272 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 272 insertions(+) create mode 100644 src/calibre/gui2/save.py diff --git a/src/calibre/gui2/save.py b/src/calibre/gui2/save.py new file mode 100644 index 0000000000..68c4055864 --- /dev/null +++ b/src/calibre/gui2/save.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python +# vim:fileencoding=utf-8 +from __future__ import (unicode_literals, division, absolute_import, + print_function) + +__license__ = 'GPL v3' +__copyright__ = '2014, Kovid Goyal ' + +import traceback, errno, os +from collections import namedtuple, defaultdict +from tempfile import SpooledTemporaryFile +from functools import partial + +from PyQt5.Qt import QObject, Qt, pyqtSignal + +from calibre.customize.ui import apply_null_metadata +from calibre.db.errors import NoSuchFormat +from calibre.ebooks.metadata import authors_to_string +from calibre.ebooks.metadata.opf2 import metadata_to_opf +from calibre.gui2 import error_dialog, warning_dialog, gprefs, open_local_file +from calibre.gui2.dialogs.progress import ProgressDialog +from calibre.utils.formatter_functions import load_user_template_functions +from calibre.library.save_to_disk import sanitize_args, get_path_components, update_metadata + +BookId = namedtuple('BookId', 'title authors') + +def ensure_unique_components(data): # {{{ + cmap = {} + for book_id, (mi, components) in data.iteritems(): + c = tuple(components) + if c in cmap: + cmap[c].add(book_id) + else: + cmap[c] = {book_id} + + for book_ids in cmap.itervalues(): + if len(book_ids) > 1: + for i, book_id in enumerate(sorted(book_ids)[1:]): + suffix = ' (%d)' % (i + 1) + components[-1] = components[-1] + suffix +# }}} + +class SpooledFile(SpooledTemporaryFile): # {{{ + + def __init__(self, file_obj, max_size=50*1024*1024): + self._file_obj = file_obj + SpooledTemporaryFile.__init__(self, max_size) + + def rollover(self): + if self._rolled: + return + orig = self._file + newfile = self._file = self._file_obj + del self._TemporaryFileArgs + + newfile.write(orig.getvalue()) + newfile.seek(orig.tell(), 0) + + self._rolled = True + + def truncate(self, *args): + # The stdlib SpooledTemporaryFile implementation of truncate() doesn't + # allow specifying a size. + self._file.truncate(*args) +# }}} + +class Saver(QObject): + + do_one_signal = pyqtSignal() + + def __init__(self, book_ids, db, opts, root, parent=None): + QObject.__init__(self, parent) + if parent is not None: + setattr(parent, 'no_gc_%s' % id(self), self) + self.db = db.new_api + self.plugboards = self.db.pref('plugboards', {}) + self.template_functions = self.db.pref('user_template_functions', []) + load_user_template_functions('', self.template_functions) + self.collected_data = {} + self.errors = defaultdict(list) + self._book_id_data = {} + self.all_book_ids = frozenset(book_ids) + self.pd = ProgressDialog(_('Saving...'), _('Collecting metadata...'), min=0, max=0, parent=parent, icon='save.png') + self.do_one_signal.connect(self.tick, type=Qt.QueuedConnection) + self.do_one = self.do_one_collect + self.ids_to_collect = iter(self.all_book_ids) + + self.pd.show() + self.root, self.opts, self.path_length = sanitize_args(root, opts) + self.do_one_signal.emit() + + def tick(self): + if self.pd.canceled: + self.pd.close() + self.break_cycles() + return + self.do_one() + + def break_cycles(self): + p = self.parent() + if p is not None: + setattr(p, 'no_gc_%s' % id(self), None) + self.plugboards = self.template_functions = self.collected_data = self.all_book_ids = self.pd = self.db = None + + def book_id_data(self, book_id): + ans = self._book_id_data.get(book_id) + if ans is None: + try: + ans = BookId(self.db.field_for('title', book_id), self.db.field_for('authors', book_id)) + except Exception: + ans = BookId((_('Unknown') + ' (%d)' % book_id), (_('Unknown'),)) + self._book_id_data[book_id] = ans + return ans + + def do_one_collect(self): + try: + book_id = next(self.ids_to_collect) + except StopIteration: + self.collection_finished() + return + try: + self.collect_data(book_id) + except Exception: + self.errors[book_id].append(('critical', traceback.format_exc())) + self.do_one_signal.emit() + + def collect_data(self, book_id): + mi = self.db.get_metadata(book_id) + components = get_path_components(self.opts, mi, book_id, self.path_length) + self.collected_data[book_id] = (mi, components) + + def collection_finished(self): + self.do_one = self.do_one_write + ensure_unique_components(self.collected_data) + self.ids_to_write = iter(self.collected_data) + self.pd.max = len(self.collected_data) + self.pd.value = 0 + self.do_one_signal.emit() + + def do_one_write(self): + try: + book_id = next(self.ids_to_write) + except StopIteration: + self.writing_finished() + return + self.pd.msg = self.book_id_data(book_id).title + self.pd.value += 1 + try: + self.fmts_to_write = self.write_book(book_id) + except Exception: + self.errors[book_id].append(('critical', traceback.format_exc())) + self.fmts_to_write = iter(()) + self.do_one = self.do_one_fmt + self.do_one_signal.emit() + + def writing_finished(self): + self.pd.close() + self.report() + self.break_cycles() + if gprefs['show_files_after_save']: + open_local_file(self.root) + + def write_book(self, book_id): + mi, components = self.collected_data[book_id] + base_path = os.path.join(self.root, *components) + base_dir = os.path.dirname(base_path) + fmts = {f.lower() for f in self.db.formats(book_id)} + if self.opts.formats != 'all': + asked_formats = {x.lower().strip() for x in self.opts.formats.split(',')} + fmts = asked_formats.intersection(fmts) + + if not fmts and not self.opts.write_opf and not self.opts.save_cover: + return + try: + os.makedirs(base_dir) + except EnvironmentError as err: + if err.errno != errno.EEXIST: + raise + + cdata = self.db.cover(book_id) + mi.cover, mi.cover_data = None, (None, None) + + if self.opts.save_cover and cdata: + with lopen(base_path + os.extsep + 'jpg', 'wb') as f: + f.write(cdata) + mi.cover = os.path.basename(f.name) + + if self.opts.write_opf: + opf = metadata_to_opf(mi) + with lopen(base_path + os.extsep + 'opf', 'wb') as f: + f.write(opf) + + mi.cover, mi.cover_data = None, (None, None) + + return ((book_id, f, base_path, mi, cdata) for f in fmts) + + def do_one_fmt(self): + try: + args = next(self.fmts_to_write) + except StopIteration: + del self.fmts_to_write + self.do_one = self.do_one_write + self.do_one_signal.emit() + return + try: + self.write_fmt(*args) + except Exception: + self.errors[args[0]].append(('fmt', (args[1], traceback.format_exc()))) + self.do_one_signal.emit() + + def report_update_metadata_error(self, book_id, fmt, tb): + self.errors[book_id].append(('metadata', (fmt, tb))) + + def write_fmt(self, book_id, fmt, base_path, mi, cdata): + fmtpath = base_path + os.extsep + fmt + written = False + with lopen(fmtpath, 'w+b') as f: + sf = SpooledTemporaryFile(f) + try: + self.db.copy_format_to(book_id, fmt, sf) + except NoSuchFormat: + pass + else: + if self.opts.update_metadata: + sf.seek(0) + with apply_null_metadata: + update_metadata(mi, fmt, sf, self.plugboards, cdata, + error_report=partial(self.report_update_metadata_error, book_id)) + sf.rollover() + written = True + if not written: + os.remove(fmtpath) + + def format_report(self): + report = [] + a = report.append + + def indent(text): + return '\xa0\xa0\xa0\xa0' + '\n\xa0\xa0\xa0\xa0'.join(text.splitlines()) + + for book_id, errors in self.errors.iteritems(): + types = {t for t, data in errors} + title, authors = self.book_id_data(book_id).title, authors_to_string(self.book_id_data(book_id).authors[:1]) + if report: + a('\n' + ('_'*80) + '\n') + if 'critical' in types: + a(_('Failed to save: {0} by {1} to disk, with error:').format(title, authors)) + for t, tb in errors: + if t == 'critical': + a(indent(tb)) + else: + errs = defaultdict(list) + for t, data in errors: + errs[t].append(data) + for fmt, tb in errs['fmt']: + a(_('Failed to save the {2} format of: {0} by {1} to disk, with error:').format(title, authors, fmt.upper())) + a(indent(tb)), a('') + for fmt, tb in errs['metadata']: + a(_('Failed to update the metadata in the {2} format of: {0} by {1} to disk, with error:').format(title, authors, fmt.upper())) + a(indent(tb)), a('') + return '\n'.join(report) + + def report(self): + if not self.errors: + return + if len(self.errors) == len(self.all_book_ids): + msg = _('Failed to save any books to disk, click "Show details" for more information') + d = error_dialog + else: + msg = _('Failed to save some books to disk, click "Show details" for more information') + d = warning_dialog + d(self.parent(), _('Error while saving'), msg, det_msg=self.format_report(), show=True)