1#!/usr/local/bin/python3.8
2# vim:fileencoding=utf-8
3
4
5__license__ = 'GPL v3'
6__copyright__ = '2014, Kovid Goyal <kovid at kovidgoyal.net>'
7
8import traceback, errno, os, time, shutil
9from collections import namedtuple, defaultdict
10
11from qt.core import QObject, Qt, pyqtSignal
12
13from calibre import prints, force_unicode
14from calibre.constants import DEBUG
15from calibre.customize.ui import can_set_metadata
16from calibre.db.errors import NoSuchFormat
17from calibre.ebooks.metadata import authors_to_string
18from calibre.ebooks.metadata.opf2 import metadata_to_opf
19from calibre.ptempfile import PersistentTemporaryDirectory, SpooledTemporaryFile
20from calibre.gui2 import error_dialog, warning_dialog, gprefs, open_local_file
21from calibre.gui2.dialogs.progress import ProgressDialog
22from calibre.utils.formatter_functions import load_user_template_functions
23from calibre.utils.ipc.pool import Pool, Failure
24from calibre.library.save_to_disk import sanitize_args, get_path_components, find_plugboard, plugboard_save_to_disk_value
25from polyglot.builtins import iteritems, itervalues
26from polyglot.queue import Empty
27
28BookId = namedtuple('BookId', 'title authors')
29
30
31def ensure_unique_components(data):  # {{{
32    cmap = defaultdict(set)
33    bid_map = {}
34    for book_id, (mi, components, fmts) in iteritems(data):
35        cmap[tuple(components)].add(book_id)
36        bid_map[book_id] = components
37
38    for book_ids in itervalues(cmap):
39        if len(book_ids) > 1:
40            for i, book_id in enumerate(sorted(book_ids)[1:]):
41                suffix = ' (%d)' % (i + 1)
42                components = bid_map[book_id]
43                components[-1] = components[-1] + suffix
44# }}}
45
46
47class SpooledFile(SpooledTemporaryFile):  # {{{
48
49    def __init__(self, file_obj, max_size=50*1024*1024):
50        self._file_obj = file_obj
51        SpooledTemporaryFile.__init__(self, max_size)
52
53    def rollover(self):
54        if self._rolled:
55            return
56        orig = self._file
57        newfile = self._file = self._file_obj
58        del self._TemporaryFileArgs
59
60        newfile.write(orig.getvalue())
61        newfile.seek(orig.tell(), 0)
62
63        self._rolled = True
64
65# }}}
66
67
68class Saver(QObject):
69
70    do_one_signal = pyqtSignal()
71
72    def __init__(self, book_ids, db, opts, root, parent=None, pool=None):
73        QObject.__init__(self, parent)
74        self.db = db.new_api
75        self.plugboards = self.db.pref('plugboards', {})
76        self.template_functions = self.db.pref('user_template_functions', [])
77        load_user_template_functions('', self.template_functions)
78        self.collected_data = {}
79        self.errors = defaultdict(list)
80        self._book_id_data = {}
81        self.all_book_ids = frozenset(book_ids)
82        self.pd = ProgressDialog(_('Saving %d books...') % len(self.all_book_ids), _('Collecting metadata...'), min=0, max=0, parent=parent, icon='save.png')
83        self.do_one_signal.connect(self.tick, type=Qt.ConnectionType.QueuedConnection)
84        self.do_one = self.do_one_collect
85        self.ids_to_collect = iter(self.all_book_ids)
86        self.tdir = PersistentTemporaryDirectory('_save_to_disk')
87        self.pool = pool
88
89        self.pd.show()
90        self.root, self.opts, self.path_length = sanitize_args(root, opts)
91        self.do_one_signal.emit()
92        if DEBUG:
93            self.start_time = time.time()
94
95    def tick(self):
96        if self.pd.canceled:
97            self.pd.close()
98            self.pd.deleteLater()
99            self.break_cycles()
100            return
101        self.do_one()
102
103    def break_cycles(self):
104        shutil.rmtree(self.tdir, ignore_errors=True)
105        if self.pool is not None:
106            self.pool.shutdown()
107        self.setParent(None)
108        self.jobs = self.pool = self.plugboards = self.template_functions = self.collected_data = self.all_book_ids = self.pd = self.db = None  # noqa
109        self.deleteLater()
110
111    def book_id_data(self, book_id):
112        ans = self._book_id_data.get(book_id)
113        if ans is None:
114            try:
115                ans = BookId(self.db.field_for('title', book_id), self.db.field_for('authors', book_id))
116            except Exception:
117                ans = BookId((_('Unknown') + ' (%d)' % book_id), (_('Unknown'),))
118            self._book_id_data[book_id] = ans
119        return ans
120
121    def do_one_collect(self):
122        try:
123            book_id = next(self.ids_to_collect)
124        except StopIteration:
125            self.collection_finished()
126            return
127        try:
128            self.collect_data(book_id)
129        except Exception:
130            self.errors[book_id].append(('critical', traceback.format_exc()))
131        self.do_one_signal.emit()
132
133    def collect_data(self, book_id):
134        mi = self.db.get_metadata(book_id)
135        self._book_id_data[book_id] = BookId(mi.title, mi.authors)
136        components = get_path_components(self.opts, mi, book_id, self.path_length)
137        self.collected_data[book_id] = (mi, components, {fmt.lower() for fmt in self.db.formats(book_id)})
138
139    def collection_finished(self):
140        self.do_one = self.do_one_write
141        ensure_unique_components(self.collected_data)
142        self.ids_to_write = iter(self.collected_data)
143        self.pd.title = _('Copying files and writing metadata...') if self.opts.update_metadata else _(
144            'Copying files...')
145        self.pd.max = len(self.collected_data)
146        self.pd.value = 0
147        if self.opts.update_metadata:
148            all_fmts = {fmt for data in itervalues(self.collected_data) for fmt in data[2]}
149            plugboards_cache = {fmt:find_plugboard(plugboard_save_to_disk_value, fmt, self.plugboards) for fmt in all_fmts}
150            self.pool = Pool(name='SaveToDisk') if self.pool is None else self.pool
151            try:
152                self.pool.set_common_data(plugboards_cache)
153            except Failure as err:
154                error_dialog(self.pd, _('Critical failure'), _(
155                    'Could not save books to disk, click "Show details" for more information'),
156                    det_msg=force_unicode(err.failure_message) + '\n' + force_unicode(err.details), show=True)
157                self.pd.canceled = True
158        self.do_one_signal.emit()
159
160    def do_one_write(self):
161        try:
162            book_id = next(self.ids_to_write)
163        except StopIteration:
164            self.writing_finished()
165            return
166        if not self.opts.update_metadata:
167            self.pd.msg = self.book_id_data(book_id).title
168            self.pd.value += 1
169        try:
170            self.write_book(book_id, *self.collected_data[book_id])
171        except Exception:
172            self.errors[book_id].append(('critical', traceback.format_exc()))
173        self.consume_results()
174        self.do_one_signal.emit()
175
176    def consume_results(self):
177        if self.pool is not None:
178            while True:
179                try:
180                    worker_result = self.pool.results.get_nowait()
181                except Empty:
182                    break
183                book_id = worker_result.id
184                if worker_result.is_terminal_failure:
185                    error_dialog(self.pd, _('Critical failure'), _(
186                        'The update metadata worker process crashed while processing'
187                        ' the book %s. Saving is aborted.') % self.book_id_data(book_id).title, show=True)
188                    self.pd.canceled = True
189                    return
190                result = worker_result.result
191                self.pd.value += 1
192                self.pd.msg = self.book_id_data(book_id).title
193                if result.err is not None:
194                    self.errors[book_id].append(('metadata', (None, result.err + '\n' + result.traceback)))
195                if result.value:
196                    for fmt, tb in result.value:
197                        self.errors[book_id].append(('metadata', (fmt, tb)))
198
199    def write_book(self, book_id, mi, components, fmts):
200        base_path = os.path.join(self.root, *components)
201        base_dir = os.path.dirname(base_path)
202        if self.opts.formats and self.opts.formats != 'all':
203            asked_formats = {x.lower().strip() for x in self.opts.formats.split(',')}
204            fmts = asked_formats.intersection(fmts)
205            if not fmts:
206                self.errors[book_id].append(('critical', _('Requested formats not available')))
207                return
208
209        if not fmts and not self.opts.write_opf and not self.opts.save_cover:
210            return
211
212        # On windows python incorrectly raises an access denied exception
213        # when trying to create the root of a drive, like C:\
214        if os.path.dirname(base_dir) != base_dir:
215            try:
216                os.makedirs(base_dir)
217            except OSError as err:
218                if err.errno != errno.EEXIST:
219                    raise
220
221        if self.opts.update_metadata:
222            d = {}
223            d['last_modified'] = mi.last_modified.isoformat()
224
225        cdata = self.db.cover(book_id)
226        mi.cover, mi.cover_data = None, (None, None)
227
228        if cdata:
229            fname = None
230            if self.opts.save_cover:
231                fname = base_path + os.extsep + 'jpg'
232                mi.cover = os.path.basename(fname)
233            elif self.opts.update_metadata:
234                fname = os.path.join(self.tdir, '%d.jpg' % book_id)
235
236            if fname:
237                with lopen(fname, 'wb') as f:
238                    f.write(cdata)
239                if self.opts.update_metadata:
240                    d['cover'] = fname
241
242        fname = None
243        if self.opts.write_opf:
244            fname = base_path + os.extsep + 'opf'
245        elif self.opts.update_metadata:
246            fname = os.path.join(self.tdir, '%d.opf' % book_id)
247        if fname:
248            opf = metadata_to_opf(mi)
249            with lopen(fname, 'wb') as f:
250                f.write(opf)
251            if self.opts.update_metadata:
252                d['opf'] = fname
253        mi.cover, mi.cover_data = None, (None, None)
254        if self.opts.update_metadata:
255            d['fmts'] = []
256        for fmt in fmts:
257            try:
258                fmtpath = self.write_fmt(book_id, fmt, base_path)
259                if fmtpath and self.opts.update_metadata and can_set_metadata(fmt):
260                    d['fmts'].append(fmtpath)
261            except Exception:
262                self.errors[book_id].append(('fmt', (fmt, traceback.format_exc())))
263        if self.opts.update_metadata:
264            if d['fmts']:
265                try:
266                    self.pool(book_id, 'calibre.library.save_to_disk', 'update_serialized_metadata', d)
267                except Failure as err:
268                    error_dialog(self.pd, _('Critical failure'), _(
269                        'Could not save books to disk, click "Show details" for more information'),
270                        det_msg=str(err.failure_message) + '\n' + str(err.details), show=True)
271                    self.pd.canceled = True
272            else:
273                self.pd.value += 1
274                self.pd.msg = self.book_id_data(book_id).title
275
276    def write_fmt(self, book_id, fmt, base_path):
277        fmtpath = base_path + os.extsep + fmt
278        written = False
279        with lopen(fmtpath, 'w+b') as f:
280            try:
281                self.db.copy_format_to(book_id, fmt, f)
282                written = True
283            except NoSuchFormat:
284                self.errors[book_id].append(('fmt', (fmt, _('No %s format file present') % fmt.upper())))
285        if not written:
286            os.remove(fmtpath)
287        if written:
288            return fmtpath
289
290    def writing_finished(self):
291        if not self.opts.update_metadata:
292            self.updating_metadata_finished()
293        else:
294            self.do_one = self.do_one_update
295            self.do_one_signal.emit()
296
297    def do_one_update(self):
298        self.consume_results()
299        try:
300            self.pool.wait_for_tasks(0.1)
301        except Failure as err:
302            error_dialog(self.pd, _('Critical failure'), _(
303                'Could not save books to disk, click "Show details" for more information'),
304                det_msg=str(err.failure_message) + '\n' + str(err.details), show=True)
305            self.pd.canceled = True
306        except RuntimeError:
307            pass  # tasks not completed
308        else:
309            self.consume_results()
310            return self.updating_metadata_finished()
311        self.do_one_signal.emit()
312
313    def updating_metadata_finished(self):
314        if DEBUG:
315            prints('Saved %d books in %.1f seconds' % (len(self.all_book_ids), time.time() - self.start_time))
316        self.pd.close()
317        self.pd.deleteLater()
318        self.report()
319        self.break_cycles()
320        if gprefs['show_files_after_save']:
321            open_local_file(self.root)
322
323    def format_report(self):
324        report = []
325        a = report.append
326
327        def indent(text):
328            text = force_unicode(text)
329            return '\xa0\xa0\xa0\xa0' + '\n\xa0\xa0\xa0\xa0'.join(text.splitlines())
330
331        for book_id, errors in iteritems(self.errors):
332            types = {t for t, data in errors}
333            title, authors = self.book_id_data(book_id).title, authors_to_string(self.book_id_data(book_id).authors[:1])
334            if report:
335                a('\n' + ('_'*70) + '\n')
336            if 'critical' in types:
337                a(_('Failed to save: {0} by {1} to disk, with error:').format(title, authors))
338                for t, tb in errors:
339                    if t == 'critical':
340                        a(indent(tb))
341            else:
342                errs = defaultdict(list)
343                for t, data in errors:
344                    errs[t].append(data)
345                for fmt, tb in errs['fmt']:
346                    a(_('Failed to save the {2} format of: {0} by {1} to disk, with error:').format(title, authors, fmt.upper()))
347                    a(indent(tb)), a('')
348                for fmt, tb in errs['metadata']:
349                    if fmt:
350                        a(_('Failed to update the metadata in the {2} format of: {0} by {1}, with error:').format(title, authors, fmt.upper()))
351                    else:
352                        a(_('Failed to update the metadata in all formats of: {0} by {1}, with error:').format(title, authors))
353                    a(indent(tb)), a('')
354        return '\n'.join(report)
355
356    def report(self):
357        if not self.errors:
358            return
359        err_types = {e[0] for errors in itervalues(self.errors) for e in errors}
360        if err_types == {'metadata'}:
361            msg = _('Failed to update metadata in some books, click "Show details" for more information')
362            d = warning_dialog
363        elif len(self.errors) == len(self.all_book_ids):
364            msg = _('Failed to save any books to disk, click "Show details" for more information')
365            d = error_dialog
366        else:
367            msg = _('Failed to save some books to disk, click "Show details" for more information')
368            d = warning_dialog
369        d(self.parent(), _('Error while saving'), msg, det_msg=self.format_report(), show=True)
370