1#!/usr/local/bin/python3.8
2# vim:fileencoding=UTF-8:ts=4:sw=4:sta:et:sts=4:fdm=marker:ai
3# License: GPLv3 Copyright: 2013, Kovid Goyal <kovid at kovidgoyal.net>
4
5
6import os
7import sys
8import tempfile
9import textwrap
10from functools import partial
11from qt.core import (
12    QAbstractItemView, QApplication, QCheckBox, QCursor, QDialog, QDialogButtonBox,
13    QEvent, QFrame, QGridLayout, QIcon, QInputDialog, QItemSelectionModel,
14    QKeySequence, QLabel, QMenu, QPushButton, QScrollArea, QSize, QSizePolicy,
15    QStackedWidget, Qt, QToolButton, QTreeWidget, QTreeWidgetItem, QVBoxLayout,
16    QWidget, pyqtSignal
17)
18from threading import Thread
19from time import monotonic
20
21from calibre.constants import TOC_DIALOG_APP_UID, islinux, iswindows
22from calibre.ebooks.oeb.polish.container import AZW3Container, get_container
23from calibre.ebooks.oeb.polish.toc import (
24    TOC, add_id, commit_toc, from_files, from_links, from_xpaths, get_toc
25)
26from calibre.gui2 import (
27    Application, error_dialog, info_dialog, question_dialog, set_app_uid
28)
29from calibre.gui2.convert.xpath_wizard import XPathEdit
30from calibre.gui2.progress_indicator import ProgressIndicator
31from calibre.gui2.toc.location import ItemEdit
32from calibre.ptempfile import reset_base_dir
33from calibre.utils.config import JSONConfig
34from calibre.utils.filenames import atomic_rename
35from calibre.utils.logging import GUILog
36
37ICON_SIZE = 24
38
39
40class XPathDialog(QDialog):  # {{{
41
42    def __init__(self, parent, prefs):
43        QDialog.__init__(self, parent)
44        self.prefs = prefs
45        self.setWindowTitle(_('Create ToC from XPath'))
46        self.l = l = QVBoxLayout()
47        self.setLayout(l)
48        self.la = la = QLabel(_(
49            'Specify a series of XPath expressions for the different levels of'
50            ' the Table of Contents. You can use the wizard buttons to help'
51            ' you create XPath expressions.'))
52        la.setWordWrap(True)
53        l.addWidget(la)
54        self.widgets = []
55        for i in range(5):
56            la = _('Level %s ToC:')%('&%d'%(i+1))
57            xp = XPathEdit(self)
58            xp.set_msg(la)
59            self.widgets.append(xp)
60            l.addWidget(xp)
61
62        self.bb = bb = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok|QDialogButtonBox.StandardButton.Cancel)
63        bb.accepted.connect(self.accept)
64        bb.rejected.connect(self.reject)
65        self.ssb = b = bb.addButton(_('&Save settings'), QDialogButtonBox.ButtonRole.ActionRole)
66        b.clicked.connect(self.save_settings)
67        self.load_button = b = bb.addButton(_('&Load settings'), QDialogButtonBox.ButtonRole.ActionRole)
68        self.load_menu = QMenu(b)
69        b.setMenu(self.load_menu)
70        self.setup_load_button()
71        self.remove_duplicates_cb = QCheckBox(_('Do not add duplicate entries at the same level'))
72        self.remove_duplicates_cb.setChecked(self.prefs.get('xpath_toc_remove_duplicates', True))
73        l.addWidget(self.remove_duplicates_cb)
74        l.addStretch()
75        l.addWidget(bb)
76        self.resize(self.sizeHint() + QSize(50, 75))
77
78    def save_settings(self):
79        xpaths = self.xpaths
80        if not xpaths:
81            return error_dialog(self, _('No XPaths'),
82                                _('No XPaths have been entered'), show=True)
83        if not self.check():
84            return
85        name, ok = QInputDialog.getText(self, _('Choose name'),
86                _('Choose a name for these settings'))
87        if ok:
88            name = str(name).strip()
89            if name:
90                saved = self.prefs.get('xpath_toc_settings', {})
91                # in JSON all keys have to be strings
92                saved[name] = {str(i):x for i, x in enumerate(xpaths)}
93                self.prefs.set('xpath_toc_settings', saved)
94                self.setup_load_button()
95
96    def setup_load_button(self):
97        saved = self.prefs.get('xpath_toc_settings', {})
98        m = self.load_menu
99        m.clear()
100        self.__actions = []
101        a = self.__actions.append
102        for name in sorted(saved):
103            a(m.addAction(name, partial(self.load_settings, name)))
104        m.addSeparator()
105        a(m.addAction(_('Remove saved settings'), self.clear_settings))
106        self.load_button.setEnabled(bool(saved))
107
108    def clear_settings(self):
109        self.prefs.set('xpath_toc_settings', {})
110        self.setup_load_button()
111
112    def load_settings(self, name):
113        saved = self.prefs.get('xpath_toc_settings', {}).get(name, {})
114        for i, w in enumerate(self.widgets):
115            txt = saved.get(str(i), '')
116            w.edit.setText(txt)
117
118    def check(self):
119        for w in self.widgets:
120            if not w.check():
121                error_dialog(self, _('Invalid XPath'),
122                    _('The XPath expression %s is not valid.')%w.xpath,
123                             show=True)
124                return False
125        return True
126
127    def accept(self):
128        if self.check():
129            self.prefs.set('xpath_toc_remove_duplicates', self.remove_duplicates_cb.isChecked())
130            super().accept()
131
132    @property
133    def xpaths(self):
134        return [w.xpath for w in self.widgets if w.xpath.strip()]
135# }}}
136
137
138class ItemView(QStackedWidget):  # {{{
139
140    add_new_item = pyqtSignal(object, object)
141    delete_item = pyqtSignal()
142    flatten_item = pyqtSignal()
143    go_to_root = pyqtSignal()
144    create_from_xpath = pyqtSignal(object, object)
145    create_from_links = pyqtSignal()
146    create_from_files = pyqtSignal()
147    flatten_toc = pyqtSignal()
148
149    def __init__(self, parent, prefs):
150        QStackedWidget.__init__(self, parent)
151        self.prefs = prefs
152        self.setMinimumWidth(250)
153        self.root_pane = rp = QWidget(self)
154        self.item_pane = ip = QWidget(self)
155        self.current_item = None
156        sa = QScrollArea(self)
157        sa.setWidgetResizable(True)
158        sa.setWidget(rp)
159        self.addWidget(sa)
160        sa = QScrollArea(self)
161        sa.setWidgetResizable(True)
162        sa.setWidget(ip)
163        self.addWidget(sa)
164
165        self.l1 = la = QLabel('<p>'+_(
166            'You can edit existing entries in the Table of Contents by clicking them'
167            ' in the panel to the left.')+'<p>'+_(
168            'Entries with a green tick next to them point to a location that has '
169            'been verified to exist. Entries with a red dot are broken and may need'
170            ' to be fixed.'))
171        la.setStyleSheet('QLabel { margin-bottom: 20px }')
172        la.setWordWrap(True)
173        l = rp.l = QVBoxLayout()
174        rp.setLayout(l)
175        l.addWidget(la)
176        self.add_new_to_root_button = b = QPushButton(_('Create a &new entry'))
177        b.clicked.connect(self.add_new_to_root)
178        l.addWidget(b)
179        l.addStretch()
180
181        self.cfmhb = b = QPushButton(_('Generate ToC from &major headings'))
182        b.clicked.connect(self.create_from_major_headings)
183        b.setToolTip(textwrap.fill(_(
184            'Generate a Table of Contents from the major headings in the book.'
185            ' This will work if the book identifies its headings using HTML'
186            ' heading tags. Uses the <h1>, <h2> and <h3> tags.')))
187        l.addWidget(b)
188        self.cfmab = b = QPushButton(_('Generate ToC from &all headings'))
189        b.clicked.connect(self.create_from_all_headings)
190        b.setToolTip(textwrap.fill(_(
191            'Generate a Table of Contents from all the headings in the book.'
192            ' This will work if the book identifies its headings using HTML'
193            ' heading tags. Uses the <h1-6> tags.')))
194        l.addWidget(b)
195
196        self.lb = b = QPushButton(_('Generate ToC from &links'))
197        b.clicked.connect(self.create_from_links)
198        b.setToolTip(textwrap.fill(_(
199            'Generate a Table of Contents from all the links in the book.'
200            ' Links that point to destinations that do not exist in the book are'
201            ' ignored. Also multiple links with the same destination or the same'
202            ' text are ignored.'
203        )))
204        l.addWidget(b)
205
206        self.cfb = b = QPushButton(_('Generate ToC from &files'))
207        b.clicked.connect(self.create_from_files)
208        b.setToolTip(textwrap.fill(_(
209            'Generate a Table of Contents from individual files in the book.'
210            ' Each entry in the ToC will point to the start of the file, the'
211            ' text of the entry will be the "first line" of text from the file.'
212        )))
213        l.addWidget(b)
214
215        self.xpb = b = QPushButton(_('Generate ToC from &XPath'))
216        b.clicked.connect(self.create_from_user_xpath)
217        b.setToolTip(textwrap.fill(_(
218            'Generate a Table of Contents from arbitrary XPath expressions.'
219        )))
220        l.addWidget(b)
221
222        self.fal = b = QPushButton(_('&Flatten the ToC'))
223        b.clicked.connect(self.flatten_toc)
224        b.setToolTip(textwrap.fill(_(
225            'Flatten the Table of Contents, putting all entries at the top level'
226        )))
227        l.addWidget(b)
228
229        l.addStretch()
230        self.w1 = la = QLabel(_('<b>WARNING:</b> calibre only supports the '
231                                'creation of linear ToCs in AZW3 files. In a '
232                                'linear ToC every entry must point to a '
233                                'location after the previous entry. If you '
234                                'create a non-linear ToC it will be '
235                                'automatically re-arranged inside the AZW3 file.'
236                            ))
237        la.setWordWrap(True)
238        l.addWidget(la)
239
240        l = ip.l = QGridLayout()
241        ip.setLayout(l)
242        la = ip.heading = QLabel('')
243        l.addWidget(la, 0, 0, 1, 2)
244        la.setWordWrap(True)
245        la = ip.la = QLabel(_(
246            'You can move this entry around the Table of Contents by drag '
247            'and drop or using the up and down buttons to the left'))
248        la.setWordWrap(True)
249        l.addWidget(la, 1, 0, 1, 2)
250
251        # Item status
252        ip.hl1 = hl =  QFrame()
253        hl.setFrameShape(QFrame.Shape.HLine)
254        l.addWidget(hl, l.rowCount(), 0, 1, 2)
255        self.icon_label = QLabel()
256        self.status_label = QLabel()
257        self.status_label.setWordWrap(True)
258        l.addWidget(self.icon_label, l.rowCount(), 0)
259        l.addWidget(self.status_label, l.rowCount()-1, 1)
260        ip.hl2 = hl =  QFrame()
261        hl.setFrameShape(QFrame.Shape.HLine)
262        l.addWidget(hl, l.rowCount(), 0, 1, 2)
263
264        # Edit/remove item
265        rs = l.rowCount()
266        ip.b1 = b = QPushButton(QIcon(I('edit_input.png')),
267            _('Change the &location this entry points to'), self)
268        b.clicked.connect(self.edit_item)
269        l.addWidget(b, l.rowCount()+1, 0, 1, 2)
270        ip.b2 = b = QPushButton(QIcon(I('trash.png')),
271            _('&Remove this entry'), self)
272        l.addWidget(b, l.rowCount(), 0, 1, 2)
273        b.clicked.connect(self.delete_item)
274        ip.hl3 = hl =  QFrame()
275        hl.setFrameShape(QFrame.Shape.HLine)
276        l.addWidget(hl, l.rowCount(), 0, 1, 2)
277        l.setRowMinimumHeight(rs, 20)
278
279        # Add new item
280        rs = l.rowCount()
281        ip.b3 = b = QPushButton(QIcon(I('plus.png')), _('New entry &inside this entry'))
282        connect_lambda(b.clicked, self, lambda self: self.add_new('inside'))
283        l.addWidget(b, l.rowCount()+1, 0, 1, 2)
284        ip.b4 = b = QPushButton(QIcon(I('plus.png')), _('New entry &above this entry'))
285        connect_lambda(b.clicked, self, lambda self: self.add_new('before'))
286        l.addWidget(b, l.rowCount(), 0, 1, 2)
287        ip.b5 = b = QPushButton(QIcon(I('plus.png')), _('New entry &below this entry'))
288        connect_lambda(b.clicked, self, lambda self: self.add_new('after'))
289        l.addWidget(b, l.rowCount(), 0, 1, 2)
290        # Flatten entry
291        ip.b3 = b = QPushButton(QIcon(I('heuristics.png')), _('&Flatten this entry'))
292        b.clicked.connect(self.flatten_item)
293        b.setToolTip(_('All children of this entry are brought to the same '
294                       'level as this entry.'))
295        l.addWidget(b, l.rowCount()+1, 0, 1, 2)
296
297        ip.hl4 = hl =  QFrame()
298        hl.setFrameShape(QFrame.Shape.HLine)
299        l.addWidget(hl, l.rowCount(), 0, 1, 2)
300        l.setRowMinimumHeight(rs, 20)
301
302        # Return to welcome
303        rs = l.rowCount()
304        ip.b4 = b = QPushButton(QIcon(I('back.png')), _('&Return to welcome screen'))
305        b.clicked.connect(self.go_to_root)
306        b.setToolTip(_('Go back to the top level view'))
307        l.addWidget(b, l.rowCount()+1, 0, 1, 2)
308
309        l.setRowMinimumHeight(rs, 20)
310
311        l.addWidget(QLabel(), l.rowCount(), 0, 1, 2)
312        l.setColumnStretch(1, 10)
313        l.setRowStretch(l.rowCount()-1, 10)
314        self.w2 = la = QLabel(self.w1.text())
315        self.w2.setWordWrap(True)
316        l.addWidget(la, l.rowCount(), 0, 1, 2)
317
318    def ask_if_duplicates_should_be_removed(self):
319        return not question_dialog(self, _('Remove duplicates'), _(
320            'Should headings with the same text at the same level be included?'),
321            yes_text=_('&Include duplicates'), no_text=_('&Remove duplicates'))
322
323    def create_from_major_headings(self):
324        self.create_from_xpath.emit(['//h:h%d'%i for i in range(1, 4)],
325                self.ask_if_duplicates_should_be_removed())
326
327    def create_from_all_headings(self):
328        self.create_from_xpath.emit(['//h:h%d'%i for i in range(1, 7)],
329                self.ask_if_duplicates_should_be_removed())
330
331    def create_from_user_xpath(self):
332        d = XPathDialog(self, self.prefs)
333        if d.exec() == QDialog.DialogCode.Accepted and d.xpaths:
334            self.create_from_xpath.emit(d.xpaths, d.remove_duplicates_cb.isChecked())
335
336    def hide_azw3_warning(self):
337        self.w1.setVisible(False), self.w2.setVisible(False)
338
339    def add_new_to_root(self):
340        self.add_new_item.emit(None, None)
341
342    def add_new(self, where):
343        self.add_new_item.emit(self.current_item, where)
344
345    def edit_item(self):
346        self.add_new_item.emit(self.current_item, None)
347
348    def __call__(self, item):
349        if item is None:
350            self.current_item = None
351            self.setCurrentIndex(0)
352        else:
353            self.current_item = item
354            self.setCurrentIndex(1)
355            self.populate_item_pane()
356
357    def populate_item_pane(self):
358        item = self.current_item
359        name = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
360        self.item_pane.heading.setText('<h2>%s</h2>'%name)
361        self.icon_label.setPixmap(item.data(0, Qt.ItemDataRole.DecorationRole
362                                            ).pixmap(32, 32))
363        tt = _('This entry points to an existing destination')
364        toc = item.data(0, Qt.ItemDataRole.UserRole)
365        if toc.dest_exists is False:
366            tt = _('The location this entry points to does not exist')
367        elif toc.dest_exists is None:
368            tt = ''
369        self.status_label.setText(tt)
370
371    def data_changed(self, item):
372        if item is self.current_item:
373            self.populate_item_pane()
374
375# }}}
376
377
378NODE_FLAGS = (Qt.ItemFlag.ItemIsDragEnabled|Qt.ItemFlag.ItemIsEditable|Qt.ItemFlag.ItemIsEnabled|Qt.ItemFlag.ItemIsSelectable|Qt.ItemFlag.ItemIsDropEnabled)
379
380
381class TreeWidget(QTreeWidget):  # {{{
382
383    edit_item = pyqtSignal()
384    history_state_changed = pyqtSignal()
385
386    def __init__(self, parent):
387        QTreeWidget.__init__(self, parent)
388        self.history = []
389        self.setHeaderLabel(_('Table of Contents'))
390        self.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
391        self.setDragEnabled(True)
392        self.setSelectionMode(QAbstractItemView.SelectionMode.ExtendedSelection)
393        self.viewport().setAcceptDrops(True)
394        self.setDropIndicatorShown(True)
395        self.setDragDropMode(QAbstractItemView.DragDropMode.InternalMove)
396        self.setAutoScroll(True)
397        self.setAutoScrollMargin(ICON_SIZE*2)
398        self.setDefaultDropAction(Qt.DropAction.MoveAction)
399        self.setAutoExpandDelay(1000)
400        self.setAnimated(True)
401        self.setMouseTracking(True)
402        self.in_drop_event = False
403        self.root = self.invisibleRootItem()
404        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
405        self.customContextMenuRequested.connect(self.show_context_menu)
406
407    def push_history(self):
408        self.history.append(self.serialize_tree())
409        self.history_state_changed.emit()
410
411    def pop_history(self):
412        if self.history:
413            self.unserialize_tree(self.history.pop())
414            self.history_state_changed.emit()
415
416    def commitData(self, editor):
417        self.push_history()
418        return QTreeWidget.commitData(self, editor)
419
420    def iter_items(self, parent=None):
421        if parent is None:
422            parent = self.invisibleRootItem()
423        for i in range(parent.childCount()):
424            child = parent.child(i)
425            yield child
426            yield from self.iter_items(parent=child)
427
428    def update_status_tip(self, item):
429        c = item.data(0, Qt.ItemDataRole.UserRole)
430        if c is not None:
431            frag = c.frag or ''
432            if frag:
433                frag = '#'+frag
434            item.setStatusTip(0, _('<b>Title</b>: {0} <b>Dest</b>: {1}{2}').format(
435                c.title, c.dest, frag))
436
437    def serialize_tree(self):
438
439        def serialize_node(node):
440            return {
441                'title': node.data(0, Qt.ItemDataRole.DisplayRole),
442                'toc_node': node.data(0, Qt.ItemDataRole.UserRole),
443                'icon': node.data(0, Qt.ItemDataRole.DecorationRole),
444                'tooltip': node.data(0, Qt.ItemDataRole.ToolTipRole),
445                'is_selected': node.isSelected(),
446                'is_expanded': node.isExpanded(),
447                'children': list(map(serialize_node, (node.child(i) for i in range(node.childCount())))),
448            }
449
450        node = self.invisibleRootItem()
451        return {'children': list(map(serialize_node, (node.child(i) for i in range(node.childCount()))))}
452
453    def unserialize_tree(self, serialized):
454
455        def unserialize_node(dict_node, parent):
456            n = QTreeWidgetItem(parent)
457            n.setData(0, Qt.ItemDataRole.DisplayRole, dict_node['title'])
458            n.setData(0, Qt.ItemDataRole.UserRole, dict_node['toc_node'])
459            n.setFlags(NODE_FLAGS)
460            n.setData(0, Qt.ItemDataRole.DecorationRole, dict_node['icon'])
461            n.setData(0, Qt.ItemDataRole.ToolTipRole, dict_node['tooltip'])
462            self.update_status_tip(n)
463            n.setExpanded(dict_node['is_expanded'])
464            n.setSelected(dict_node['is_selected'])
465            for c in dict_node['children']:
466                unserialize_node(c, n)
467
468        i = self.invisibleRootItem()
469        i.takeChildren()
470        for child in serialized['children']:
471            unserialize_node(child, i)
472
473    def dropEvent(self, event):
474        self.in_drop_event = True
475        self.push_history()
476        try:
477            super().dropEvent(event)
478        finally:
479            self.in_drop_event = False
480
481    def selectedIndexes(self):
482        ans = super().selectedIndexes()
483        if self.in_drop_event:
484            # For order to be be preserved when moving by drag and drop, we
485            # have to ensure that selectedIndexes returns an ordered list of
486            # indexes.
487            sort_map = {self.indexFromItem(item):i for i, item in enumerate(self.iter_items())}
488            ans = sorted(ans, key=lambda x:sort_map.get(x, -1))
489        return ans
490
491    def highlight_item(self, item):
492        self.setCurrentItem(item, 0, QItemSelectionModel.SelectionFlag.ClearAndSelect)
493        self.scrollToItem(item)
494
495    def check_multi_selection(self):
496        if len(self.selectedItems()) > 1:
497            info_dialog(self, _('Multiple items selected'), _(
498                'You are trying to move multiple items at once, this is not supported. Instead use'
499                ' Drag and Drop to move multiple items'), show=True)
500            return False
501        return True
502
503    def move_left(self):
504        if not self.check_multi_selection():
505            return
506        self.push_history()
507        item = self.currentItem()
508        if item is not None:
509            parent = item.parent()
510            if parent is not None:
511                is_expanded = item.isExpanded() or item.childCount() == 0
512                gp = parent.parent() or self.invisibleRootItem()
513                idx = gp.indexOfChild(parent)
514                for gc in [parent.child(i) for i in range(parent.indexOfChild(item)+1, parent.childCount())]:
515                    parent.removeChild(gc)
516                    item.addChild(gc)
517                parent.removeChild(item)
518                gp.insertChild(idx+1, item)
519                if is_expanded:
520                    self.expandItem(item)
521                self.highlight_item(item)
522
523    def move_right(self):
524        if not self.check_multi_selection():
525            return
526        self.push_history()
527        item = self.currentItem()
528        if item is not None:
529            parent = item.parent() or self.invisibleRootItem()
530            idx = parent.indexOfChild(item)
531            if idx > 0:
532                is_expanded = item.isExpanded()
533                np = parent.child(idx-1)
534                parent.removeChild(item)
535                np.addChild(item)
536                if is_expanded:
537                    self.expandItem(item)
538                self.highlight_item(item)
539
540    def move_down(self):
541        if not self.check_multi_selection():
542            return
543        self.push_history()
544        item = self.currentItem()
545        if item is None:
546            if self.root.childCount() == 0:
547                return
548            item = self.root.child(0)
549            self.highlight_item(item)
550            return
551        parent = item.parent() or self.root
552        idx = parent.indexOfChild(item)
553        if idx == parent.childCount() - 1:
554            # At end of parent, need to become sibling of parent
555            if parent is self.root:
556                return
557            gp = parent.parent() or self.root
558            parent.removeChild(item)
559            gp.insertChild(gp.indexOfChild(parent)+1, item)
560        else:
561            sibling = parent.child(idx+1)
562            parent.removeChild(item)
563            sibling.insertChild(0, item)
564        self.highlight_item(item)
565
566    def move_up(self):
567        if not self.check_multi_selection():
568            return
569        self.push_history()
570        item = self.currentItem()
571        if item is None:
572            if self.root.childCount() == 0:
573                return
574            item = self.root.child(self.root.childCount()-1)
575            self.highlight_item(item)
576            return
577        parent = item.parent() or self.root
578        idx = parent.indexOfChild(item)
579        if idx == 0:
580            # At end of parent, need to become sibling of parent
581            if parent is self.root:
582                return
583            gp = parent.parent() or self.root
584            parent.removeChild(item)
585            gp.insertChild(gp.indexOfChild(parent), item)
586        else:
587            sibling = parent.child(idx-1)
588            parent.removeChild(item)
589            sibling.addChild(item)
590        self.highlight_item(item)
591
592    def del_items(self):
593        self.push_history()
594        for item in self.selectedItems():
595            p = item.parent() or self.root
596            p.removeChild(item)
597
598    def title_case(self):
599        self.push_history()
600        from calibre.utils.titlecase import titlecase
601        for item in self.selectedItems():
602            t = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
603            item.setData(0, Qt.ItemDataRole.DisplayRole, titlecase(t))
604
605    def upper_case(self):
606        self.push_history()
607        for item in self.selectedItems():
608            t = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
609            item.setData(0, Qt.ItemDataRole.DisplayRole, icu_upper(t))
610
611    def lower_case(self):
612        self.push_history()
613        for item in self.selectedItems():
614            t = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
615            item.setData(0, Qt.ItemDataRole.DisplayRole, icu_lower(t))
616
617    def swap_case(self):
618        self.push_history()
619        from calibre.utils.icu import swapcase
620        for item in self.selectedItems():
621            t = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
622            item.setData(0, Qt.ItemDataRole.DisplayRole, swapcase(t))
623
624    def capitalize(self):
625        self.push_history()
626        from calibre.utils.icu import capitalize
627        for item in self.selectedItems():
628            t = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
629            item.setData(0, Qt.ItemDataRole.DisplayRole, capitalize(t))
630
631    def bulk_rename(self):
632        from calibre.gui2.tweak_book.file_list import get_bulk_rename_settings
633        sort_map = {id(item):i for i, item in enumerate(self.iter_items())}
634        items = sorted(self.selectedItems(), key=lambda x:sort_map.get(id(x), -1))
635        settings = get_bulk_rename_settings(self, len(items), prefix=_('Chapter '), msg=_(
636            'All selected items will be renamed to the form prefix-number'), sanitize=lambda x:x, leading_zeros=False)
637        fmt, num = settings['prefix'], settings['start']
638        if fmt is not None and num is not None:
639            self.push_history()
640            for i, item in enumerate(items):
641                item.setData(0, Qt.ItemDataRole.DisplayRole, fmt % (num + i))
642
643    def keyPressEvent(self, ev):
644        if ev.key() == Qt.Key.Key_Left and ev.modifiers() & Qt.KeyboardModifier.ControlModifier:
645            self.move_left()
646            ev.accept()
647        elif ev.key() == Qt.Key.Key_Right and ev.modifiers() & Qt.KeyboardModifier.ControlModifier:
648            self.move_right()
649            ev.accept()
650        elif ev.key() == Qt.Key.Key_Up and (ev.modifiers() & Qt.KeyboardModifier.ControlModifier or ev.modifiers() & Qt.KeyboardModifier.AltModifier):
651            self.move_up()
652            ev.accept()
653        elif ev.key() == Qt.Key.Key_Down and (ev.modifiers() & Qt.KeyboardModifier.ControlModifier or ev.modifiers() & Qt.KeyboardModifier.AltModifier):
654            self.move_down()
655            ev.accept()
656        elif ev.key() in (Qt.Key.Key_Delete, Qt.Key.Key_Backspace):
657            self.del_items()
658            ev.accept()
659        else:
660            return super().keyPressEvent(ev)
661
662    def show_context_menu(self, point):
663        item = self.currentItem()
664
665        def key(k):
666            sc = str(QKeySequence(k | Qt.KeyboardModifier.ControlModifier).toString(QKeySequence.SequenceFormat.NativeText))
667            return ' [%s]'%sc
668
669        if item is not None:
670            m = QMenu(self)
671            m.addAction(QIcon(I('edit_input.png')), _('Change the location this entry points to'), self.edit_item)
672            m.addAction(QIcon(I('modified.png')), _('Bulk rename all selected items'), self.bulk_rename)
673            m.addAction(QIcon(I('trash.png')), _('Remove all selected items'), self.del_items)
674            m.addSeparator()
675            ci = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
676            p = item.parent() or self.invisibleRootItem()
677            idx = p.indexOfChild(item)
678            if idx > 0:
679                m.addAction(QIcon(I('arrow-up.png')), (_('Move "%s" up')%ci)+key(Qt.Key.Key_Up), self.move_up)
680            if idx + 1 < p.childCount():
681                m.addAction(QIcon(I('arrow-down.png')), (_('Move "%s" down')%ci)+key(Qt.Key.Key_Down), self.move_down)
682            if item.parent() is not None:
683                m.addAction(QIcon(I('back.png')), (_('Unindent "%s"')%ci)+key(Qt.Key.Key_Left), self.move_left)
684            if idx > 0:
685                m.addAction(QIcon(I('forward.png')), (_('Indent "%s"')%ci)+key(Qt.Key.Key_Right), self.move_right)
686
687            m.addSeparator()
688            case_menu = QMenu(_('Change case'), m)
689            case_menu.addAction(_('Upper case'), self.upper_case)
690            case_menu.addAction(_('Lower case'), self.lower_case)
691            case_menu.addAction(_('Swap case'), self.swap_case)
692            case_menu.addAction(_('Title case'), self.title_case)
693            case_menu.addAction(_('Capitalize'), self.capitalize)
694            m.addMenu(case_menu)
695
696            m.exec(QCursor.pos())
697# }}}
698
699
700class TOCView(QWidget):  # {{{
701
702    add_new_item = pyqtSignal(object, object)
703
704    def __init__(self, parent, prefs):
705        QWidget.__init__(self, parent)
706        self.toc_title = None
707        self.prefs = prefs
708        l = self.l = QGridLayout()
709        self.setLayout(l)
710        self.tocw = t = TreeWidget(self)
711        self.tocw.edit_item.connect(self.edit_item)
712        l.addWidget(t, 0, 0, 7, 3)
713        self.up_button = b = QToolButton(self)
714        b.setIcon(QIcon(I('arrow-up.png')))
715        b.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
716        l.addWidget(b, 0, 3)
717        b.setToolTip(_('Move current entry up [Ctrl+Up]'))
718        b.clicked.connect(self.move_up)
719
720        self.left_button = b = QToolButton(self)
721        b.setIcon(QIcon(I('back.png')))
722        b.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
723        l.addWidget(b, 2, 3)
724        b.setToolTip(_('Unindent the current entry [Ctrl+Left]'))
725        b.clicked.connect(self.tocw.move_left)
726
727        self.del_button = b = QToolButton(self)
728        b.setIcon(QIcon(I('trash.png')))
729        b.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
730        l.addWidget(b, 3, 3)
731        b.setToolTip(_('Remove all selected entries'))
732        b.clicked.connect(self.del_items)
733
734        self.right_button = b = QToolButton(self)
735        b.setIcon(QIcon(I('forward.png')))
736        b.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
737        l.addWidget(b, 4, 3)
738        b.setToolTip(_('Indent the current entry [Ctrl+Right]'))
739        b.clicked.connect(self.tocw.move_right)
740
741        self.down_button = b = QToolButton(self)
742        b.setIcon(QIcon(I('arrow-down.png')))
743        b.setIconSize(QSize(ICON_SIZE, ICON_SIZE))
744        l.addWidget(b, 6, 3)
745        b.setToolTip(_('Move current entry down [Ctrl+Down]'))
746        b.clicked.connect(self.move_down)
747        self.expand_all_button = b = QPushButton(_('&Expand all'))
748        col = 7
749        l.addWidget(b, col, 0)
750        b.clicked.connect(self.tocw.expandAll)
751        self.collapse_all_button = b = QPushButton(_('&Collapse all'))
752        b.clicked.connect(self.tocw.collapseAll)
753        l.addWidget(b, col, 1)
754        self.default_msg = _('Double click on an entry to change the text')
755        self.hl = hl = QLabel(self.default_msg)
756        hl.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Ignored)
757        l.addWidget(hl, col, 2, 1, -1)
758        self.item_view = i = ItemView(self, self.prefs)
759        self.item_view.delete_item.connect(self.delete_current_item)
760        i.add_new_item.connect(self.add_new_item)
761        i.create_from_xpath.connect(self.create_from_xpath)
762        i.create_from_links.connect(self.create_from_links)
763        i.create_from_files.connect(self.create_from_files)
764        i.flatten_item.connect(self.flatten_item)
765        i.flatten_toc.connect(self.flatten_toc)
766        i.go_to_root.connect(self.go_to_root)
767        l.addWidget(i, 0, 4, col, 1)
768
769        l.setColumnStretch(2, 10)
770
771    def edit_item(self):
772        self.item_view.edit_item()
773
774    def event(self, e):
775        if e.type() == QEvent.Type.StatusTip:
776            txt = str(e.tip()) or self.default_msg
777            self.hl.setText(txt)
778        return super().event(e)
779
780    def item_title(self, item):
781        return str(item.data(0, Qt.ItemDataRole.DisplayRole) or '')
782
783    def del_items(self):
784        self.tocw.del_items()
785
786    def delete_current_item(self):
787        item = self.tocw.currentItem()
788        if item is not None:
789            self.tocw.push_history()
790            p = item.parent() or self.root
791            p.removeChild(item)
792
793    def iter_items(self, parent=None):
794        yield from self.tocw.iter_items(parent=parent)
795
796    def flatten_toc(self):
797        self.tocw.push_history()
798        found = True
799        while found:
800            found = False
801            for item in self.iter_items():
802                if item.childCount() > 0:
803                    self._flatten_item(item)
804                    found = True
805                    break
806
807    def flatten_item(self):
808        self.tocw.push_history()
809        self._flatten_item(self.tocw.currentItem())
810
811    def _flatten_item(self, item):
812        if item is not None:
813            p = item.parent() or self.root
814            idx = p.indexOfChild(item)
815            children = [item.child(i) for i in range(item.childCount())]
816            for child in reversed(children):
817                item.removeChild(child)
818                p.insertChild(idx+1, child)
819
820    def go_to_root(self):
821        self.tocw.setCurrentItem(None)
822
823    def highlight_item(self, item):
824        self.tocw.highlight_item(item)
825
826    def move_up(self):
827        self.tocw.move_up()
828
829    def move_down(self):
830        self.tocw.move_down()
831
832    def data_changed(self, top_left, bottom_right):
833        for r in range(top_left.row(), bottom_right.row()+1):
834            idx = self.tocw.model().index(r, 0, top_left.parent())
835            new_title = str(idx.data(Qt.ItemDataRole.DisplayRole) or '').strip()
836            toc = idx.data(Qt.ItemDataRole.UserRole)
837            if toc is not None:
838                toc.title = new_title or _('(Untitled)')
839            item = self.tocw.itemFromIndex(idx)
840            self.tocw.update_status_tip(item)
841            self.item_view.data_changed(item)
842
843    def create_item(self, parent, child, idx=-1):
844        if idx == -1:
845            c = QTreeWidgetItem(parent)
846        else:
847            c = QTreeWidgetItem()
848            parent.insertChild(idx, c)
849        self.populate_item(c, child)
850        return c
851
852    def populate_item(self, c, child):
853        c.setData(0, Qt.ItemDataRole.DisplayRole, child.title or _('(Untitled)'))
854        c.setData(0, Qt.ItemDataRole.UserRole, child)
855        c.setFlags(NODE_FLAGS)
856        c.setData(0, Qt.ItemDataRole.DecorationRole, self.icon_map[child.dest_exists])
857        if child.dest_exists is False:
858            c.setData(0, Qt.ItemDataRole.ToolTipRole, _(
859                'The location this entry point to does not exist:\n%s')
860                %child.dest_error)
861        else:
862            c.setData(0, Qt.ItemDataRole.ToolTipRole, None)
863
864        self.tocw.update_status_tip(c)
865
866    def __call__(self, ebook):
867        self.ebook = ebook
868        if not isinstance(ebook, AZW3Container):
869            self.item_view.hide_azw3_warning()
870        self.toc = get_toc(self.ebook)
871        self.toc_lang, self.toc_uid = self.toc.lang, self.toc.uid
872        self.toc_title = self.toc.toc_title
873        self.blank = QIcon(I('blank.png'))
874        self.ok = QIcon(I('ok.png'))
875        self.err = QIcon(I('dot_red.png'))
876        self.icon_map = {None:self.blank, True:self.ok, False:self.err}
877
878        def process_item(toc_node, parent):
879            for child in toc_node:
880                c = self.create_item(parent, child)
881                process_item(child, c)
882
883        root = self.root = self.tocw.invisibleRootItem()
884        root.setData(0, Qt.ItemDataRole.UserRole, self.toc)
885        process_item(self.toc, root)
886        self.tocw.model().dataChanged.connect(self.data_changed)
887        self.tocw.currentItemChanged.connect(self.current_item_changed)
888        self.tocw.setCurrentItem(None)
889
890    def current_item_changed(self, current, previous):
891        self.item_view(current)
892
893    def update_item(self, item, where, name, frag, title):
894        if isinstance(frag, tuple):
895            frag = add_id(self.ebook, name, *frag)
896        child = TOC(title, name, frag)
897        child.dest_exists = True
898        self.tocw.push_history()
899        if item is None:
900            # New entry at root level
901            c = self.create_item(self.root, child)
902            self.tocw.setCurrentItem(c, 0, QItemSelectionModel.SelectionFlag.ClearAndSelect)
903            self.tocw.scrollToItem(c)
904        else:
905            if where is None:
906                # Editing existing entry
907                self.populate_item(item, child)
908            else:
909                if where == 'inside':
910                    parent = item
911                    idx = -1
912                else:
913                    parent = item.parent() or self.root
914                    idx = parent.indexOfChild(item)
915                    if where == 'after':
916                        idx += 1
917                c = self.create_item(parent, child, idx=idx)
918                self.tocw.setCurrentItem(c, 0, QItemSelectionModel.SelectionFlag.ClearAndSelect)
919                self.tocw.scrollToItem(c)
920
921    def create_toc(self):
922        root = TOC()
923
924        def process_node(parent, toc_parent):
925            for i in range(parent.childCount()):
926                item = parent.child(i)
927                title = str(item.data(0, Qt.ItemDataRole.DisplayRole) or '').strip()
928                toc = item.data(0, Qt.ItemDataRole.UserRole)
929                dest, frag = toc.dest, toc.frag
930                toc = toc_parent.add(title, dest, frag)
931                process_node(item, toc)
932
933        process_node(self.tocw.invisibleRootItem(), root)
934        return root
935
936    def insert_toc_fragment(self, toc):
937
938        def process_node(root, tocparent, added):
939            for child in tocparent:
940                item = self.create_item(root, child)
941                added.append(item)
942                process_node(item, child, added)
943
944        self.tocw.push_history()
945        nodes = []
946        process_node(self.root, toc, nodes)
947        self.highlight_item(nodes[0])
948
949    def create_from_xpath(self, xpaths, remove_duplicates=True):
950        toc = from_xpaths(self.ebook, xpaths)
951        if len(toc) == 0:
952            return error_dialog(self, _('No items found'),
953                _('No items were found that could be added to the Table of Contents.'), show=True)
954        if remove_duplicates:
955            toc.remove_duplicates()
956        self.insert_toc_fragment(toc)
957
958    def create_from_links(self):
959        toc = from_links(self.ebook)
960        if len(toc) == 0:
961            return error_dialog(self, _('No items found'),
962                _('No links were found that could be added to the Table of Contents.'), show=True)
963        self.insert_toc_fragment(toc)
964
965    def create_from_files(self):
966        toc = from_files(self.ebook)
967        if len(toc) == 0:
968            return error_dialog(self, _('No items found'),
969                _('No files were found that could be added to the Table of Contents.'), show=True)
970        self.insert_toc_fragment(toc)
971
972    def undo(self):
973        self.tocw.pop_history()
974
975
976# }}}
977
978
979te_prefs = JSONConfig('toc-editor')
980
981
982class TOCEditor(QDialog):  # {{{
983
984    explode_done = pyqtSignal(object)
985    writing_done = pyqtSignal(object)
986
987    def __init__(self, pathtobook, title=None, parent=None, prefs=None, write_result_to=None):
988        QDialog.__init__(self, parent)
989        self.last_reject_at = self.last_accept_at = -1000
990        self.write_result_to = write_result_to
991        self.prefs = prefs or te_prefs
992        self.pathtobook = pathtobook
993        self.working = True
994
995        t = title or os.path.basename(pathtobook)
996        self.book_title = t
997        self.setWindowTitle(_('Edit the ToC in %s')%t)
998        self.setWindowIcon(QIcon(I('highlight_only_on.png')))
999
1000        l = self.l = QVBoxLayout()
1001        self.setLayout(l)
1002
1003        self.stacks = s = QStackedWidget(self)
1004        l.addWidget(s)
1005        self.loading_widget = lw = QWidget(self)
1006        s.addWidget(lw)
1007        ll = self.ll = QVBoxLayout()
1008        lw.setLayout(ll)
1009        self.pi = pi = ProgressIndicator()
1010        pi.setDisplaySize(QSize(200, 200))
1011        pi.startAnimation()
1012        ll.addWidget(pi, alignment=Qt.AlignmentFlag.AlignHCenter|Qt.AlignmentFlag.AlignCenter)
1013        la = self.wait_label = QLabel(_('Loading %s, please wait...')%t)
1014        la.setWordWrap(True)
1015        f = la.font()
1016        f.setPointSize(20), la.setFont(f)
1017        ll.addWidget(la, alignment=Qt.AlignmentFlag.AlignHCenter|Qt.AlignmentFlag.AlignTop)
1018        self.toc_view = TOCView(self, self.prefs)
1019        self.toc_view.add_new_item.connect(self.add_new_item)
1020        self.toc_view.tocw.history_state_changed.connect(self.update_history_buttons)
1021        s.addWidget(self.toc_view)
1022        self.item_edit = ItemEdit(self)
1023        s.addWidget(self.item_edit)
1024
1025        bb = self.bb = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok|QDialogButtonBox.StandardButton.Cancel)
1026        l.addWidget(bb)
1027        bb.accepted.connect(self.accept)
1028        bb.rejected.connect(self.reject)
1029        self.undo_button = b = bb.addButton(_('&Undo'), QDialogButtonBox.ButtonRole.ActionRole)
1030        b.setToolTip(_('Undo the last action, if any'))
1031        b.setIcon(QIcon(I('edit-undo.png')))
1032        b.clicked.connect(self.toc_view.undo)
1033
1034        self.explode_done.connect(self.read_toc, type=Qt.ConnectionType.QueuedConnection)
1035        self.writing_done.connect(self.really_accept, type=Qt.ConnectionType.QueuedConnection)
1036
1037        r = self.screen().availableSize()
1038        self.resize(r.width() - 100, r.height() - 100)
1039        geom = self.prefs.get('toc_editor_window_geom', None)
1040        if geom is not None:
1041            QApplication.instance().safe_restore_geometry(self, bytes(geom))
1042        self.stacks.currentChanged.connect(self.update_history_buttons)
1043        self.update_history_buttons()
1044
1045    def update_history_buttons(self):
1046        self.undo_button.setVisible(self.stacks.currentIndex() == 1)
1047        self.undo_button.setEnabled(bool(self.toc_view.tocw.history))
1048
1049    def add_new_item(self, item, where):
1050        self.item_edit(item, where)
1051        self.stacks.setCurrentIndex(2)
1052
1053    def accept(self):
1054        if monotonic() - self.last_accept_at < 1:
1055            return
1056        self.last_accept_at = monotonic()
1057        if self.stacks.currentIndex() == 2:
1058            self.toc_view.update_item(*self.item_edit.result)
1059            self.prefs['toc_edit_splitter_state'] = bytearray(self.item_edit.splitter.saveState())
1060            self.stacks.setCurrentIndex(1)
1061        elif self.stacks.currentIndex() == 1:
1062            self.working = False
1063            Thread(target=self.write_toc).start()
1064            self.pi.startAnimation()
1065            self.wait_label.setText(_('Writing %s, please wait...')%
1066                                    self.book_title)
1067            self.stacks.setCurrentIndex(0)
1068            self.bb.setEnabled(False)
1069
1070    def really_accept(self, tb):
1071        self.prefs['toc_editor_window_geom'] = bytearray(self.saveGeometry())
1072        if tb:
1073            error_dialog(self, _('Failed to write book'),
1074                _('Could not write %s. Click "Show details" for'
1075                  ' more information.')%self.book_title, det_msg=tb, show=True)
1076            super().reject()
1077            return
1078        self.write_result(0)
1079        super().accept()
1080
1081    def reject(self):
1082        if not self.bb.isEnabled():
1083            return
1084        if monotonic() - self.last_reject_at < 1:
1085            return
1086        self.last_reject_at = monotonic()
1087        if self.stacks.currentIndex() == 2:
1088            self.prefs['toc_edit_splitter_state'] = bytearray(self.item_edit.splitter.saveState())
1089            self.stacks.setCurrentIndex(1)
1090        else:
1091            self.working = False
1092            self.prefs['toc_editor_window_geom'] = bytearray(self.saveGeometry())
1093            self.write_result(1)
1094            super().reject()
1095
1096    def write_result(self, res):
1097        if self.write_result_to:
1098            with tempfile.NamedTemporaryFile(dir=os.path.dirname(self.write_result_to), delete=False) as f:
1099                src = f.name
1100                f.write(str(res).encode('utf-8'))
1101                f.flush()
1102            atomic_rename(src, self.write_result_to)
1103
1104    def start(self):
1105        t = Thread(target=self.explode)
1106        t.daemon = True
1107        self.log = GUILog()
1108        t.start()
1109
1110    def explode(self):
1111        tb = None
1112        try:
1113            self.ebook = get_container(self.pathtobook, log=self.log)
1114        except:
1115            import traceback
1116            tb = traceback.format_exc()
1117        if self.working:
1118            self.working = False
1119            self.explode_done.emit(tb)
1120
1121    def read_toc(self, tb):
1122        if tb:
1123            error_dialog(self, _('Failed to load book'),
1124                _('Could not load %s. Click "Show details" for'
1125                  ' more information.')%self.book_title, det_msg=tb, show=True)
1126            self.reject()
1127            return
1128        self.pi.stopAnimation()
1129        self.toc_view(self.ebook)
1130        self.item_edit.load(self.ebook)
1131        self.stacks.setCurrentIndex(1)
1132
1133    def write_toc(self):
1134        tb = None
1135        try:
1136            toc = self.toc_view.create_toc()
1137            toc.toc_title = getattr(self.toc_view, 'toc_title', None)
1138            commit_toc(self.ebook, toc, lang=self.toc_view.toc_lang,
1139                    uid=self.toc_view.toc_uid)
1140            self.ebook.commit()
1141        except:
1142            import traceback
1143            tb = traceback.format_exc()
1144        self.writing_done.emit(tb)
1145
1146# }}}
1147
1148
1149def main(path=None, title=None):
1150    # Ensure we can continue to function if GUI is closed
1151    os.environ.pop('CALIBRE_WORKER_TEMP_DIR', None)
1152    reset_base_dir()
1153    if iswindows:
1154        # Ensure that all instances are grouped together in the task bar. This
1155        # prevents them from being grouped with viewer/editor process when
1156        # launched from within calibre, as both use calibre-parallel.exe
1157        set_app_uid(TOC_DIALOG_APP_UID)
1158
1159    with open(path + '.started', 'w'):
1160        pass
1161    override = 'calibre-gui' if islinux else None
1162    app = Application([], override_program_name=override)
1163    d = TOCEditor(path, title=title, write_result_to=path + '.result')
1164    d.start()
1165    ret = 1
1166    if d.exec() == QDialog.DialogCode.Accepted:
1167        ret = 0
1168    del d
1169    del app
1170    raise SystemExit(ret)
1171
1172
1173if __name__ == '__main__':
1174    main(path=sys.argv[-1], title='test')
1175    os.remove(sys.argv[-1] + '.lock')
1176