Projet

Général

Profil

Anomalie #290 » menu_builder_dialog.py

Leslie Lemaire, 30/06/2020 11:56

 
1
# -*- coding: utf-8 -*-
2
"""
3
MenuBuilder - Create your own menus with your favorite layers
4

    
5
copyright            : (C) 2015 by Oslandia
6
email                : infos@oslandia.com
7

    
8
/***************************************************************************
9
 *                                                                         *
10
 *   This program is free software; you can redistribute it and/or modify  *
11
 *   it under the terms of the GNU General Public License as published by  *
12
 *   the Free Software Foundation; either version 2 of the License, or     *
13
 *   (at your option) any later version.                                   *
14
 *                                                                         *
15
 ***************************************************************************/
16
"""
17
import os
18
import re
19
import json
20
from contextlib import contextmanager
21
from collections import defaultdict
22
from functools import wraps, partial
23

    
24
import psycopg2
25

    
26
from PyQt5.QtCore import Qt, QRect, QSortFilterProxyModel
27
from PyQt5.QtWidgets import (
28
    QAction, QMessageBox, QDialog, QMenu, QTreeView,
29
    QAbstractItemView, QDockWidget, QWidget, QVBoxLayout,
30
    QSizePolicy, QLineEdit, QDialogButtonBox
31
)
32

    
33
from PyQt5.QtGui import (
34
    QIcon, QStandardItem,
35
    QStandardItemModel
36
)
37
from qgis.core import (
38
    QgsProject, QgsBrowserModel, QgsDataSourceUri, QgsSettings,
39
    QgsCredentials, QgsVectorLayer, QgsMimeDataUtils, QgsRasterLayer
40
)
41

    
42
from .menu_builder_dialog_base import Ui_Dialog
43

    
44
QGIS_MIMETYPE = 'application/x-vnd.qgis.qgis.uri'
45

    
46

    
47
ICON_MAPPER = {
48
    'postgres': ":/plugins/MenuBuilder/resources/postgis.svg",
49
    'WMS': ":/plugins/MenuBuilder/resources/wms.svg",
50
    'WFS': ":/plugins/MenuBuilder/resources/wfs.svg",
51
    'OWS': ":/plugins/MenuBuilder/resources/ows.svg",
52
    'spatialite': ":/plugins/MenuBuilder/resources/spatialite.svg",
53
    'mssql': ":/plugins/MenuBuilder/resources/mssql.svg",
54
    'gdal': ":/plugins/MenuBuilder/resources/gdal.svg",
55
    'ogr': ":/plugins/MenuBuilder/resources/ogr.svg",
56
}
57

    
58

    
59
class MenuBuilderDialog(QDialog, Ui_Dialog):
60

    
61
    def __init__(self, uiparent):
62
        super().__init__()
63

    
64
        self.setupUi(self)
65

    
66
        # reference to caller
67
        self.uiparent = uiparent
68

    
69
        self.combo_profile.lineEdit().setPlaceholderText(self.tr("Profile name"))
70

    
71
        # add icons
72
        self.button_add_menu.setIcon(QIcon(":/plugins/MenuBuilder/resources/plus.svg"))
73
        self.button_delete_profile.setIcon(QIcon(":/plugins/MenuBuilder/resources/delete.svg"))
74

    
75
        # custom qtreeview
76
        self.target = CustomQtTreeView(self)
77
        self.target.setGeometry(QRect(440, 150, 371, 451))
78
        self.target.setAcceptDrops(True)
79
        self.target.setDragEnabled(True)
80
        self.target.setDragDropMode(QAbstractItemView.DragDrop)
81
        self.target.setObjectName("target")
82
        self.target.setDropIndicatorShown(True)
83
        self.target.setSelectionMode(QAbstractItemView.ExtendedSelection)
84
        self.target.setHeaderHidden(True)
85
        sizePolicy = QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
86
        sizePolicy.setHorizontalStretch(0)
87
        sizePolicy.setVerticalStretch(0)
88
        sizePolicy.setHeightForWidth(self.target.sizePolicy().hasHeightForWidth())
89
        self.target.setSizePolicy(sizePolicy)
90
        self.target.setAutoFillBackground(True)
91
        self.verticalLayout_2.addWidget(self.target)
92
        self.menumodel = MenuTreeModel(self)
93
        self.target.setModel(self.menumodel)
94
        self.target.setAnimated(True)
95
        self.target.setDefaultDropAction(Qt.MoveAction)
96

    
97
        self.browser = QgsBrowserModel()
98
        self.browser.initialize()
99
        self.source.setModel(self.browser)
100
        self.source.setHeaderHidden(True)
101
        self.source.setDragEnabled(True)
102
        self.source.setSelectionMode(QAbstractItemView.ExtendedSelection)
103

    
104
        # add a dock widget
105
        self.dock_widget = QDockWidget("Menus")
106
        self.dock_widget.resize(400, 300)
107
        self.dock_widget.setFloating(True)
108
        self.dock_widget.setObjectName(self.tr("Menu Tree"))
109
        self.dock_widget_content = QWidget()
110
        self.dock_widget.setWidget(self.dock_widget_content)
111
        dock_layout = QVBoxLayout()
112
        self.dock_widget_content.setLayout(dock_layout)
113
        self.dock_view = DockQtTreeView(self.dock_widget_content)
114
        self.dock_view.setDragDropMode(QAbstractItemView.DragOnly)
115
        self.dock_menu_filter = QLineEdit()
116
        self.dock_menu_filter.setPlaceholderText(self.tr("Filter by table description (postgis only)"))
117
        dock_layout.addWidget(self.dock_menu_filter)
118
        dock_layout.addWidget(self.dock_view)
119
        self.dock_view.setHeaderHidden(True)
120
        self.dock_view.setDragEnabled(True)
121
        self.dock_view.setSelectionMode(QAbstractItemView.ExtendedSelection)
122
        self.dock_view.setAnimated(True)
123
        self.dock_view.setObjectName("treeView")
124
        self.proxy_model = LeafFilterProxyModel(self)
125
        self.proxy_model.setFilterRole(Qt.ToolTipRole)
126
        self.proxy_model.setFilterCaseSensitivity(Qt.CaseInsensitive)
127

    
128
        self.profile_list = []
129
        self.table = 'qgis_menubuilder_metadata'
130

    
131
        self.layer_handler = {
132
            'vector': self.load_vector,
133
            'raster': self.load_raster
134
        }
135

    
136
        # connect signals and handlers
137
        self.combo_database.activated.connect(partial(self.set_connection, dbname=None))
138
        self.combo_schema.activated.connect(self.update_profile_list)
139
        self.combo_profile.activated.connect(partial(self.update_model_idx, self.menumodel))
140
        self.button_add_menu.released.connect(self.add_menu)
141
        self.button_delete_profile.released.connect(self.delete_profile)
142
        self.dock_menu_filter.textEdited.connect(self.filter_update)
143
        self.dock_view.doubleClicked.connect(self.load_from_index)
144

    
145
        self.buttonBox.rejected.connect(self.reject)
146
        self.buttonBox.accepted.connect(self.accept)
147
        self.buttonBox.button(QDialogButtonBox.Apply).clicked.connect(self.apply)
148

    
149
    def filter_update(self):
150
        text = self.dock_menu_filter.displayText()
151
        self.proxy_model.setFilterRegExp(text)
152

    
153
    def show_dock(self, state, profile=None, schema=None):
154
        if not state:
155
            # just hide widget
156
            self.dock_widget.setVisible(state)
157
            return
158
        # dock must be read only and deepcopy of model is not supported (c++ inside!)
159
        self.dock_model = MenuTreeModel(self)
160
        if profile:
161
            # bypass combobox
162
            self.update_model(self.dock_model, schema, profile)
163
        else:
164
            self.update_model_idx(self.dock_model, self.combo_profile.currentIndex())
165
        self.dock_model.setHorizontalHeaderLabels(["Menus"])
166
        self.dock_view.setEditTriggers(QAbstractItemView.NoEditTriggers)
167
        self.proxy_model.setSourceModel(self.dock_model)
168
        self.dock_view.setModel(self.proxy_model)
169
        self.dock_widget.setVisible(state)
170

    
171
    def show_menus(self, state, profile=None, schema=None):
172
        if state:
173
            self.load_menus(profile=profile, schema=schema)
174
            return
175
        # remove menus
176
        for menu in self.uiparent.menus:
177
            self.uiparent.iface.mainWindow().menuBar().removeAction(menu.menuAction())
178

    
179
    def add_menu(self):
180
        """
181
        Add a menu inside qtreeview
182
        """
183
        item = QStandardItem('NewMenu')
184
        item.setIcon(QIcon(':/plugins/MenuBuilder/resources/menu.svg'))
185
        # select current index selected and insert as a sibling
186
        brother = self.target.selectedIndexes()
187

    
188
        if not brother or not brother[0].parent():
189
            # no selection, add menu at the top level
190
            self.menumodel.insertRow(self.menumodel.rowCount(), item)
191
            return
192

    
193
        parent = self.menumodel.itemFromIndex(brother[0].parent())
194
        if not parent:
195
            self.menumodel.insertRow(self.menumodel.rowCount(), item)
196
            return
197
        parent.appendRow(item)
198

    
199
    def update_database_list(self):
200
        """update list of defined postgres connections"""
201
        settings = QgsSettings()
202
        settings.beginGroup("/PostgreSQL/connections")
203
        keys = settings.childGroups()
204
        self.combo_database.clear()
205
        self.combo_schema.clear()
206
        self.menumodel.clear()
207
        self.combo_database.addItems(keys)
208
        self.combo_database.setCurrentIndex(-1)
209
        settings.endGroup()
210
        # clear profile list
211
        self.combo_profile.clear()
212
        self.combo_profile.setCurrentIndex(-1)
213

    
214
    def set_connection(self, databaseidx, dbname=None):
215
        """
216
        Connect to selected postgresql database
217
        """
218
        selected = self.combo_database.itemText(databaseidx) or dbname
219
        if not selected:
220
            return
221

    
222
        settings = QgsSettings()
223
        settings.beginGroup("/PostgreSQL/connections/{}".format(selected))
224

    
225
        if not settings.contains("database"):
226
            # no entry?
227
            QMessageBox.critical(self, "Error", "There is no defined database connection")
228
            return
229

    
230
        uri = QgsDataSourceUri()
231

    
232
        settingsList = ["service", "host", "port", "database", "username", "password"]
233
        service, host, port, database, username, password = map(
234
            lambda x: settings.value(x, "", type=str), settingsList)
235

    
236
        useEstimatedMetadata = settings.value("estimatedMetadata", False, type=bool)
237
        sslmode = settings.enumValue("sslmode", uri.SslPrefer)
238

    
239
        settings.endGroup()
240

    
241
        if service:
242
            uri.setConnection(service, database, username, password, sslmode)
243
        else:
244
            uri.setConnection(host, port, database, username, password, sslmode)
245

    
246
        uri.setUseEstimatedMetadata(useEstimatedMetadata)
247

    
248
        # connect to db
249
        try:
250
            self.connect_to_uri(uri)
251
        except self.pg_error_types():
252
            QMessageBox.warning(
253
                self,
254
                "Plugin MenuBuilder: Message",
255
                self.tr("The database containing Menu's configuration is unavailable"),
256
                QMessageBox.Ok,
257
            )
258
            # connection not available
259
            return False
260

    
261
        # update schema list
262
        self.update_schema_list()
263
        return True
264

    
265
    @contextmanager
266
    def transaction(self):
267
        try:
268
            yield
269
            self.connection.commit()
270
        except self.pg_error_types() as e:
271
            self.connection.rollback()
272
            raise e
273

    
274
    def check_connected(func):
275
        """
276
        Decorator that checks if a database connection is active before executing function
277
        """
278
        @wraps(func)
279
        def wrapped(inst, *args, **kwargs):
280
            if not getattr(inst, 'connection', False):
281
                QMessageBox(
282
                    QMessageBox.Warning,
283
                    "Menu Builder",
284
                    inst.tr("Not connected to any database, please select one"),
285
                    QMessageBox.Ok,
286
                    inst
287
                )
288
                return
289
            if inst.connection.closed:
290
                QMessageBox(
291
                    QMessageBox.Warning,
292
                    "Menu Builder",
293
                    inst.tr("Not connected to any database, please select one"),
294
                    QMessageBox.Ok,
295
                    inst
296
                )
297
                return
298
            return func(inst, *args, **kwargs)
299
        return wrapped
300

    
301
    def connect_to_uri(self, uri):
302
        self.close_connection()
303
        self.host = uri.host() or os.environ.get('PGHOST')
304
        self.port = uri.port() or os.environ.get('PGPORT')
305

    
306
        username = uri.username() or os.environ.get('PGUSER') or os.environ.get('USER')
307
        password = uri.password() or os.environ.get('PGPASSWORD')
308
        
309
        conninfo = uri.connectionInfo()
310

    
311
        while True:
312
            try:
313
                self.connection = psycopg2.connect(uri.connectionInfo(), application_name="QGIS:MenuBuilder")
314
                break
315
            except self.pg_error_types() as e:
316
                err = str(e) or "Erreur d'authentification. Vérifiez les informations saisies."
317
                                
318
                ok, username, password = QgsCredentials.instance().get(
319
                    conninfo, username, password, err)
320
                if not ok:
321
                    raise e
322
                if username:
323
                    uri.setUsername(username)
324
                if password:
325
                    uri.setPassword(password)
326
        
327
        QgsCredentials.instance().put(conninfo, username, password) 
328
        
329
        self.pgencoding = self.connection.encoding
330

    
331
        return True
332

    
333
    def pg_error_types(self):
334
        return (
335
            psycopg2.InterfaceError,
336
            psycopg2.OperationalError,
337
            psycopg2.ProgrammingError
338
        )
339

    
340
    @check_connected
341
    def update_schema_list(self):
342
        self.combo_schema.clear()
343
        with self.transaction():
344
            cur = self.connection.cursor()
345
            cur.execute("""
346
                select nspname
347
                from pg_namespace
348
                where nspname not ilike 'pg_%'
349
                and nspname not in ('pg_catalog', 'information_schema')
350
                """)
351
            schemas = [row[0] for row in cur.fetchall()]
352
            self.combo_schema.addItems(schemas)
353

    
354
    @check_connected
355
    def update_profile_list(self, schemaidx):
356
        """
357
        update profile list from database
358
        """
359
        schema = self.combo_schema.itemText(schemaidx)
360
        with self.transaction():
361
            cur = self.connection.cursor()
362
            cur.execute("""
363
                select 1
364
                from pg_tables
365
                    where schemaname = '{0}'
366
                    and tablename = '{1}'
367
                union
368
                select 1
369
                from pg_matviews
370
                    where schemaname = '{0}'
371
                    and matviewname = '{1}'
372
                union
373
                select 1
374
                from pg_views
375
                    where schemaname = '{0}'
376
                    and viewname = '{1}'
377
                """.format(schema, self.table))
378
            tables = cur.fetchone()
379
            if not tables:
380
                box = QMessageBox(
381
                    QMessageBox.Warning,
382
                    "Menu Builder",
383
                    self.tr("Table '{}.{}' not found in this database, "
384
                            "would you like to create it now ?")
385
                    .format(schema, self.table),
386
                    QMessageBox.Cancel | QMessageBox.Yes,
387
                    self
388
                )
389
                ret = box.exec_()
390
                if ret == QMessageBox.Cancel:
391
                    return False
392
                elif ret == QMessageBox.Yes:
393
                    cur.execute("""
394
                        create table {}.{} (
395
                            id serial,
396
                            name varchar,
397
                            profile varchar,
398
                            model_index varchar,
399
                            datasource_uri text
400
                        )
401
                        """.format(schema, self.table))
402
                    self.connection.commit()
403
                    return False
404

    
405
            cur.execute("""
406
                select distinct(profile) from {}.{}
407
                """.format(schema, self.table))
408
            profiles = [row[0] for row in cur.fetchall()]
409
            saved_profile = self.combo_profile.currentText()
410
            self.combo_profile.clear()
411
            self.combo_profile.addItems(profiles)
412
            self.combo_profile.setCurrentIndex(self.combo_profile.findText(saved_profile))
413

    
414
    @check_connected
415
    def delete_profile(self):
416
        """
417
        Delete profile currently selected
418
        """
419
        idx = self.combo_profile.currentIndex()
420
        schema = self.combo_schema.currentText()
421
        profile = self.combo_profile.itemText(idx)
422
        box = QMessageBox(
423
            QMessageBox.Warning,
424
            "Menu Builder",
425
            self.tr("Delete '{}' profile ?").format(profile),
426
            QMessageBox.Cancel | QMessageBox.Yes,
427
            self
428
        )
429
        ret = box.exec_()
430
        if ret == QMessageBox.Cancel:
431
            return False
432
        elif ret == QMessageBox.Yes:
433
            self.combo_profile.removeItem(idx)
434
            with self.transaction():
435
                cur = self.connection.cursor()
436
                cur.execute("""
437
                    delete from {}.{}
438
                    where profile = '{}'
439
                    """.format(schema, self.table, profile))
440
        self.menumodel.clear()
441
        self.combo_profile.setCurrentIndex(-1)
442

    
443
    def update_model_idx(self, model, profile_index):
444
        """
445
        wrapper that checks combobox
446
        """
447
        profile = self.combo_profile.itemText(profile_index)
448
        schema = self.combo_schema.currentText()
449
        self.update_model(model, schema, profile)
450

    
451
    def sortby_modelindex(self, rows):
452
        return sorted(
453
            rows,
454
            key=lambda line: '/'.join(
455
                ['{:04}'.format(elem[0]) for elem in json.loads(line[2])]
456
            ))
457

    
458
    @check_connected
459
    def update_model(self, model, schema, profile):
460
        """
461
        Update the model by retrieving the profile given in database
462
        """
463
        menudict = {}
464

    
465
        with self.transaction():
466
            cur = self.connection.cursor()
467
            select = """
468
                select name, profile, model_index, table_comment, datasource_uri
469
                from {}.{}
470
                where profile = '{}'
471
                """.format(schema, self.table, profile)
472
            cur.execute(select)
473
            rows = cur.fetchall()
474
            model.clear()
475
            for name, profile, model_index, comment, datasource_uri in self.sortby_modelindex(rows):
476
                menu = model.invisibleRootItem()
477
                indexes = json.loads(model_index)
478
                parent = ''
479
                for idx, subname in indexes[:-1]:
480
                    parent += '{}-{}/'.format(idx, subname)
481
                    if parent in menudict:
482
                        # already created entry
483
                        menu = menudict[parent]
484
                        continue
485
                    # create menu
486
                    item = QStandardItem(subname)
487
                    uri_struct = QgsMimeDataUtils.Uri(datasource_uri)
488
                    item.setData(uri_struct)
489
                    item.setIcon(QIcon(':/plugins/MenuBuilder/resources/menu.svg'))
490
                    item.setFlags(Qt.ItemIsSelectable | Qt.ItemIsUserCheckable |
491
                                  Qt.ItemIsEnabled | Qt.ItemIsDropEnabled |
492
                                  Qt.ItemIsEditable)
493
                    item.setWhatsThis("menu")
494
                    menu.appendRow(item)
495
                    menudict[parent] = item
496
                    # set current menu to the new created item
497
                    menu = item
498

    
499
                # add leaf (layer item)
500
                item = QStandardItem(name)
501
                uri_struct = QgsMimeDataUtils.Uri(datasource_uri)
502
                # fix layer name instead of table name
503
                # usefull when the layer has been renamed in menu
504
                uri_struct.name = name
505
                if uri_struct.providerKey in ICON_MAPPER:
506
                    item.setIcon(QIcon(ICON_MAPPER[uri_struct.providerKey]))
507
                item.setData(uri_struct)
508
                # avoid placing dragged layers on it
509
                item.setDropEnabled(False)
510
                if uri_struct.providerKey == 'postgres':
511
                    # set tooltip to postgres comment
512
                    item.setToolTip(comment)
513
                menudict[parent].appendRow(item)
514

    
515
    @check_connected
516
    def save_changes(self, save_to_db=True):
517
        """
518
        Save changes in the postgres table
519
        """
520
        schema = self.combo_schema.currentText()
521
        profile = self.combo_profile.currentText()
522
        if not profile:
523
            QMessageBox(
524
                QMessageBox.Warning,
525
                "Menu Builder",
526
                self.tr("Profile cannot be empty"),
527
                QMessageBox.Ok,
528
                self
529
            ).exec_()
530
            return False
531

    
532
        if save_to_db:
533
            try:
534
                with self.transaction():
535
                    cur = self.connection.cursor()
536
                    cur.execute("delete from {}.{} where profile = '{}'".format(
537
                        schema, self.table, profile))
538
                    for item, data in self.target.iteritems():
539
                        if not data:
540
                            continue
541
                        cur.execute("""
542
                        insert into {}.{} (name,profile,model_index,datasource_uri)
543
                        values (%s, %s, %s, %s)
544
                        """.format(schema, self.table), (
545
                            item[-1][1],
546
                            profile,
547
                            json.dumps(item),
548
                            data.data())
549
                        )
550
            except Exception as exc:
551
                QMessageBox(
552
                    QMessageBox.Warning,
553
                    "Menu Builder",
554
                    exc.message.decode(self.pgencoding),
555
                    QMessageBox.Ok,
556
                    self
557
                ).exec_()
558
                return False
559

    
560
        self.save_session(
561
            self.combo_database.currentText(),
562
            schema,
563
            profile,
564
            self.activate_dock.isChecked(),
565
            self.activate_menubar.isChecked()
566
        )
567
        self.update_profile_list(self.combo_schema.currentIndex())
568
        self.show_dock(self.activate_dock.isChecked())
569
        self.show_menus(self.activate_menubar.isChecked())
570
        return True
571

    
572
    @check_connected
573
    def load_menus(self, profile=None, schema=None):
574
        """
575
        Load menus in the main windows qgis bar
576
        """
577
        if not schema:
578
            schema = self.combo_schema.currentText()
579
        if not profile:
580
            profile = self.combo_profile.currentText()
581
        # remove previous menus
582
        for menu in self.uiparent.menus:
583
            self.uiparent.iface.mainWindow().menuBar().removeAction(menu.menuAction())
584

    
585
        with self.transaction():
586
            cur = self.connection.cursor()
587
            select = """
588
                select name, profile, model_index, table_comment, datasource_uri
589
                from {}.{}
590
                where profile = '{}'
591
                """.format(schema, self.table, profile)
592
            cur.execute(select)
593
            rows = cur.fetchall()
594
        # item accessor ex: '0-menu/0-submenu/1-item/'
595
        menudict = {}
596
        # reference to parent item
597
        parent = ''
598
        # reference to qgis main menu bar
599
        menubar = self.uiparent.iface.mainWindow().menuBar()
600

    
601
        for name, profile, model_index, comment, datasource_uri in self.sortby_modelindex(rows):
602
            uri_struct = QgsMimeDataUtils.Uri(datasource_uri)
603
            indexes = json.loads(model_index)
604
            # root menu
605
            parent = '{}-{}/'.format(indexes[0][0], indexes[0][1])
606
            if parent not in menudict:
607
                menu = QMenu(self.uiparent.iface.mainWindow())
608
                self.uiparent.menus.append(menu)
609
                menu.setObjectName(indexes[0][1])
610
                menu.setTitle(indexes[0][1])
611
                menubar.insertMenu(
612
                    self.uiparent.iface.firstRightStandardMenu().menuAction(),
613
                    menu)
614
                menudict[parent] = menu
615
            else:
616
                # menu already there
617
                menu = menudict[parent]
618

    
619
            for idx, subname in indexes[1:-1]:
620
                # intermediate submenus
621
                parent += '{}-{}/'.format(idx, subname)
622
                if parent not in menudict:
623
                    submenu = menu.addMenu(subname)
624
                    submenu.setObjectName(subname)
625
                    submenu.setTitle(subname)
626
                    menu = submenu
627
                    # store it for later use
628
                    menudict[parent] = menu
629
                    continue
630
                # already treated
631
                menu = menudict[parent]
632

    
633
            # last item = layer
634
            layer = QAction(name, self)
635

    
636
            if uri_struct.providerKey in ICON_MAPPER:
637
                layer.setIcon(QIcon(ICON_MAPPER[uri_struct.providerKey]))
638

    
639
            if uri_struct.providerKey == 'postgres':
640
                # set tooltip to postgres comment
641
                layer.setStatusTip(comment)
642
                layer.setToolTip(comment)
643

    
644
            layer.setData(uri_struct.uri)
645
            layer.setWhatsThis(uri_struct.providerKey)
646
            layer.triggered.connect(self.layer_handler[uri_struct.layerType])
647
            menu.addAction(layer)  
648

    
649

    
650
    def load_from_index(self, index):
651
        """Load layers from selected item index"""
652
        item = self.dock_model.itemFromIndex(self.proxy_model.mapToSource(index))
653
        if item.whatsThis() == 'menu':
654
            return
655
        if item.data().layerType == 'vector':
656
            layer = QgsVectorLayer(
657
                item.data().uri,  # uri
658
                item.text(),  # layer name
659
                item.data().providerKey  # provider name
660
            )
661
        elif item.data().layerType == 'raster':
662
            layer = QgsRasterLayer(
663
                item.data().uri,  # uri
664
                item.text(),  # layer name
665
                item.data().providerKey  # provider name
666
            )
667
        if not layer:
668
            return
669
        QgsProject.instance().addMapLayer(layer)
670

    
671
    def load_vector(self):
672
        action = self.sender()
673

    
674
        layer = QgsVectorLayer(
675
            action.data(),  # uri
676
            action.text(),  # layer name
677
            action.whatsThis()  # provider name
678
        )
679
        QgsProject.instance().addMapLayer(layer)
680

    
681
    def load_raster(self):
682
        action = self.sender()
683
        layer = QgsRasterLayer(
684
            action.data(),  # uri
685
            action.text(),  # layer name
686
            action.whatsThis()  # provider name
687
        )
688
        QgsProject.instance().addMapLayer(layer)
689

    
690
    def accept(self):
691
        if self.save_changes():
692
            QDialog.reject(self)
693
            self.close_connection()
694

    
695
    def apply(self):
696
        if self.save_changes(save_to_db=False):
697
            QDialog.reject(self)
698
            self.close_connection()
699

    
700
    def reject(self):
701
        self.close_connection()
702
        QDialog.reject(self)
703

    
704
    def close_connection(self):
705
        """close current pg connection if exists"""
706
        if getattr(self, 'connection', False):
707
            if self.connection.closed:
708
                return
709
            self.connection.close()
710

    
711
    def save_session(self, database, schema, profile, dock, menubar):
712
        """save current profile for next session"""
713
        settings = QgsSettings()
714
        settings.setValue("MenuBuilder/database", database)
715
        settings.setValue("MenuBuilder/schema", schema)
716
        settings.setValue("MenuBuilder/profile", profile)
717
        settings.setValue("MenuBuilder/dock", dock)
718
        settings.setValue("MenuBuilder/menubar", menubar)
719

    
720
    def restore_session(self):
721
        settings = QgsSettings()
722
        database = settings.value("MenuBuilder/database", False)
723
        schema = settings.value("MenuBuilder/schema", 'public')
724
        profile = settings.value("MenuBuilder/profile", False)
725
        dock = settings.value("MenuBuilder/dock", False)
726
        menubar = settings.value("MenuBuilder/menubar", False)
727
        if not any([database, profile]):
728
            return
729

    
730
        connected = self.set_connection(0, dbname=database)
731
        if not connected:
732
            # don't try to continue
733
            return
734

    
735
        self.show_dock(bool(dock), profile=profile, schema=schema)
736
        if bool(dock):
737
            self.uiparent.iface.addDockWidget(Qt.LeftDockWidgetArea, self.dock_widget)
738
        self.show_menus(bool(menubar), profile=profile, schema=schema)
739
        self.close_connection()  
740

    
741

    
742
class CustomQtTreeView(QTreeView):
743

    
744
    def dragEnterEvent(self, event):
745
        if not event.mimeData():
746
            # don't drag menu entry
747
            return False
748
        # refuse if it's not a qgis mimetype
749
        if event.mimeData().hasFormat(QGIS_MIMETYPE):
750
            event.acceptProposedAction()
751

    
752
    def keyPressEvent(self, event):
753
        if event.key() == Qt.Key_Delete:
754
            model = self.selectionModel().model()
755
            parents = defaultdict(list)
756
            for idx in self.selectedIndexes():
757
                parents[idx.parent()].append(idx)
758
            for parent, idx_list in parents.items():
759
                for diff, index in enumerate(idx_list):
760
                    model.removeRow(index.row() - diff, parent)
761
        elif event.key() == Qt.Key_Return:
762
            pass
763
        else:
764
            super().keyPressEvent(event)
765

    
766
    def iteritems(self, level=0):
767
        """
768
        Dump model to store in database.
769
        Generates each level recursively
770
        """
771
        rowcount = self.model().rowCount()
772
        for itemidx in range(rowcount):
773
            # iterate over parents
774
            parent = self.model().itemFromIndex(self.model().index(itemidx, 0))
775
            for item, uri in self.traverse_tree(parent, []):
776
                yield item, uri
777

    
778
    def traverse_tree(self, parent, identifier):
779
        """
780
        Iterate over childs, recursively
781
        """
782
        identifier.append([parent.row(), parent.text()])
783
        for row in range(parent.rowCount()):
784
            child = parent.child(row)
785
            if child.hasChildren():
786
                # child is a menu ?
787
                for item in self.traverse_tree(child, identifier):
788
                    yield item
789
                identifier.pop()
790
            else:
791
                # add leaf
792
                sibling = list(identifier)
793
                sibling.append([child.row(), child.text()])
794
                yield sibling, child.data()
795

    
796

    
797
class DockQtTreeView(CustomQtTreeView):
798

    
799
    def keyPressEvent(self, event):
800
        """override keyevent to avoid deletion of items in the dock"""
801
        pass
802

    
803

    
804
class MenuTreeModel(QStandardItemModel):
805

    
806
    def dropMimeData(self, mimedata, action, row, column, parentIndex):
807
        """
808
        Handles the dropping of an item onto the model.
809
        De-serializes the data and inserts it into the model.
810
        """
811
        # decode data using qgis helpers
812
        uri_list = QgsMimeDataUtils.decodeUriList(mimedata)
813
        if not uri_list:
814
            return False
815
        # find parent item
816
        parent_item = self.itemFromIndex(parentIndex)
817
        if not parent_item:
818
            return False
819

    
820
        items = []
821
        for uri in uri_list:
822
            item = QStandardItem(uri.name)
823
            item.setData(uri)
824
            # avoid placing dragged layers on it
825
            item.setDropEnabled(False)
826
            if uri.providerKey in ICON_MAPPER:
827
                item.setIcon(QIcon(ICON_MAPPER[uri.providerKey]))
828
            items.append(item)
829

    
830
        if row == -1:
831
            # dropped on a Menu
832
            # add as a child at the end
833
            parent_item.appendRows(items)
834
        else:
835
            # add items at the separator shown
836
            parent_item.insertRows(row, items)
837

    
838
        return True
839

    
840
    def mimeData(self, indexes):
841
        """
842
        Used to serialize data
843
        """
844
        if not indexes:
845
            return 0
846
        items = [self.itemFromIndex(idx) for idx in indexes]
847
        if not items:
848
            return 0
849
        if not all(it.data() for it in items):
850
            return 0
851
        # reencode items
852
        mimedata = QgsMimeDataUtils.encodeUriList([item.data() for item in items])
853
        return mimedata
854

    
855
    def mimeTypes(self):
856
        return [QGIS_MIMETYPE]
857

    
858
    def supportedDropActions(self):
859
        return Qt.CopyAction | Qt.MoveAction
860

    
861

    
862
class LeafFilterProxyModel(QSortFilterProxyModel):
863
    """
864
    Class to override the following behaviour:
865
        If a parent item doesn't match the filter,
866
        none of its children will be shown.
867

    
868
    This Model matches items which are descendants
869
    or ascendants of matching items.
870
    """
871

    
872
    def filterAcceptsRow(self, row_num, source_parent):
873
        """Overriding the parent function"""
874

    
875
        # Check if the current row matches
876
        if self.filter_accepts_row_itself(row_num, source_parent):
877
            return True
878

    
879
        # Traverse up all the way to root and check if any of them match
880
        if self.filter_accepts_any_parent(source_parent):
881
            return True
882

    
883
        # Finally, check if any of the children match
884
        return self.has_accepted_children(row_num, source_parent)
885

    
886
    def filter_accepts_row_itself(self, row_num, parent):
887
        return super(LeafFilterProxyModel, self).filterAcceptsRow(row_num, parent)
888

    
889
    def filter_accepts_any_parent(self, parent):
890
        """
891
        Traverse to the root node and check if any of the
892
        ancestors match the filter
893
        """
894
        while parent.isValid():
895
            if self.filter_accepts_row_itself(parent.row(), parent.parent()):
896
                return True
897
            parent = parent.parent()
898
        return False
899

    
900
    def has_accepted_children(self, row_num, parent):
901
        """
902
        Starting from the current node as root, traverse all
903
        the descendants and test if any of the children match
904
        """
905
        model = self.sourceModel()
906
        source_index = model.index(row_num, 0, parent)
907

    
908
        children_count = model.rowCount(source_index)
909
        for i in range(children_count):
910
            if self.filterAcceptsRow(i, source_index):
911
                return True
912
        return False
    (1-1/1)