FEAT: added eurostat index browser

gdementen · gdementen · commit f0c2fc3204c2 · 2023-08-29T12:21:47.000+02:00
diff --git a/larray_editor/editor.py b/larray_editor/editor.py
@@ -1,6 +1,7 @@
 import io
 import os
 import re
+from datetime import datetime
 import sys
 from collections.abc import Sequence
 from contextlib import redirect_stdout
@@ -26,6 +27,7 @@
 import matplotlib
 import matplotlib.axes
 import numpy as np
+import pandas as pd
 
 import larray as la
 
@@ -34,12 +36,14 @@
                                  get_versions, get_documentation_url, urls, RecentlyUsedList)
 from larray_editor.arraywidget import ArrayEditorWidget
 from larray_editor.commands import EditSessionArrayCommand, EditCurrentArrayCommand
+from larray_editor.treemodel import SimpleTreeNode, SimpleLazyTreeModel
 
 from qtpy.QtCore import Qt, QUrl, QSettings
 from qtpy.QtGui import QDesktopServices, QKeySequence
 from qtpy.QtWidgets import (QMainWindow, QWidget, QListWidget, QListWidgetItem, QSplitter, QFileDialog, QPushButton,
                             QDialogButtonBox, QShortcut, QVBoxLayout, QGridLayout, QLineEdit,
-                            QCheckBox, QComboBox, QMessageBox, QDialog, QInputDialog, QLabel, QGroupBox, QRadioButton)
+                            QCheckBox, QComboBox, QMessageBox, QDialog, QInputDialog, QLabel, QGroupBox, QRadioButton,
+                            QTreeView)
 
 try:
     from qtpy.QtWidgets import QUndoStack
@@ -83,6 +87,87 @@
 DISPLAY_IN_GRID = (la.Array, np.ndarray)
 
 
+def num_leading_spaces(s):
+    i = 0
+    while s[i] == ' ':
+        i += 1
+    return i
+
+
+def indented_df_to_treenode(df, indent=4, indented_col=0, colnames=None, header=None):
+    if colnames is None:
+        colnames = df.columns.tolist()
+    if header is None:
+        header = [name.capitalize() for name in colnames]
+
+    root = SimpleTreeNode(None, header)
+    df = df[colnames]
+    parent_per_level = {0: root}
+    # iterating on the df rows directly (via df.iterrows()) is too slow
+    for row in df.values:
+        row_data = row.tolist()
+        indented_col_value = row_data[indented_col]
+        level = num_leading_spaces(indented_col_value) // indent
+        # remove indentation
+        row_data[indented_col] = indented_col_value.strip()
+
+        parent_node = parent_per_level[level]
+        node = SimpleTreeNode(parent_node, row_data)
+        parent_node.children.append(node)
+        parent_per_level[level + 1] = node
+    return root
+
+
+class EurostatBrowserDialog(QDialog):
+    def __init__(self, index, parent=None):
+        super(EurostatBrowserDialog, self).__init__(parent)
+
+        assert isinstance(index, pd.DataFrame)
+
+        # drop unused/redundant "type" column
+        index = index.drop('type', axis=1)
+
+        # display dates using locale
+        # import locale
+        # locale.setlocale(locale.LC_ALL, '')
+        # datetime.date.strftime('%x')
+
+        # CHECK: this builds the whole (model) tree in memory eagerly, so it is not lazy despite using a
+        #        Simple*Lazy*TreeModel, but it is fast enough for now. *If* it ever becomes a problem,
+        #        we could make this lazy pretty easily (see treemodel.LazyDictTreeNode for an example).
+        root = indented_df_to_treenode(index)
+        model = SimpleLazyTreeModel(root)
+        tree = QTreeView()
+        tree.setModel(model)
+        tree.setUniformRowHeights(True)
+        tree.selectionModel().currentChanged.connect(self.view_eurostat_indicator)
+        tree.setColumnWidth(0, 320)
+
+        self.resize(450, 600)
+        self.setWindowTitle("Select dataset")
+
+        # set the layout
+        layout = QVBoxLayout()
+        # layout.addWidget(toolbar)
+        layout.addWidget(tree)
+        self.setLayout(layout)
+
+    def view_eurostat_indicator(self, index):
+        from larray_eurostat import eurostat_get
+
+        node = index.internalPointer()
+        title, code, last_update_of_data, last_table_structure_change, data_start, data_end = node.data
+        if not node.children:
+            last_update_of_data = datetime.strptime(last_update_of_data, "%d.%m.%Y")
+            last_table_structure_change = datetime.strptime(last_table_structure_change, "%d.%m.%Y")
+            last_change = max(last_update_of_data, last_table_structure_change)
+            try:
+                arr = eurostat_get(code, maxage=last_change, cache_dir='__array_cache__')
+            except Exception:
+                QMessageBox.critical(self, "Error", "Failed to load {}".format(code))
+            self.parent().view_expr(arr, expr=code)
+
+
 class AbstractEditor(QMainWindow):
     """Abstract Editor Window"""
 
@@ -553,6 +638,7 @@ def _setup_file_menu(self, menu_bar):
         # ============= #
         file_menu.addSeparator()
         file_menu.addAction(create_action(self, _('&Load Example Dataset'), triggered=self.load_example))
+        file_menu.addAction(create_action(self, _('&Browse Eurostat Datasets'), triggered=self.browse_eurostat))
         # ============= #
         #    SCRIPTS    #
         # ============= #
@@ -675,7 +761,7 @@ def line_edit_update(self):
 
     def view_expr(self, array, expr):
         self._listwidget.clearSelection()
-        self.set_current_array(array, expr)
+        self.set_current_array(array, name=expr)
 
     def _display_in_grid(self, k, v):
         return not k.startswith('__') and isinstance(v, DISPLAY_IN_GRID)
@@ -1153,6 +1239,16 @@ def load_example(self):
                 filepath = AVAILABLE_EXAMPLE_DATA[dataset_name]
                 self._open_file(filepath)
 
+    def browse_eurostat(self):
+        from larray_eurostat import get_index
+        try:
+            df = get_index(cache_dir='__array_cache__', maxage=None)
+        except:
+            QMessageBox.critical(self, "Error", "Failed to fetch Eurostat dataset index")
+            return
+        dialog = EurostatBrowserDialog(df, parent=self)
+        dialog.show()
+
 
 class ArrayEditor(AbstractEditor):
     """Array Editor Dialog"""
diff --git a/larray_editor/treemodel.py b/larray_editor/treemodel.py
@@ -0,0 +1,71 @@
+from qtpy.QtCore import Qt, QModelIndex, QAbstractItemModel
+
+
+class SimpleTreeNode(object):
+    __slots__ = ['parent', 'data', 'children']
+
+    def __init__(self, parent, data):
+        self.parent = parent
+        self.data = data
+        self.children = []
+
+
+class SimpleLazyTreeModel(QAbstractItemModel):
+    def __init__(self, root, parent=None):
+        super(SimpleLazyTreeModel, self).__init__(parent)
+        assert isinstance(root, SimpleTreeNode)
+        self.root = root
+
+    def columnCount(self, index):
+        node = index.internalPointer() if index.isValid() else self.root
+        return len(node.data)
+
+    def data(self, index, role):
+        if not index.isValid():
+            return None
+
+        if role != Qt.DisplayRole:
+            return None
+
+        node = index.internalPointer()
+        return node.data[index.column()]
+
+    def flags(self, index):
+        if not index.isValid():
+            return Qt.NoItemFlags
+
+        return Qt.ItemIsEnabled | Qt.ItemIsSelectable
+
+    def headerData(self, section, orientation, role):
+        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
+            return self.root.data[section]
+        return None
+
+    def index(self, row, column, parent_index):
+        if not self.hasIndex(row, column, parent_index):
+            return QModelIndex()
+
+        parent_node = parent_index.internalPointer() if parent_index.isValid() else self.root
+        child_node = parent_node.children[row]
+        return self.createIndex(row, column, child_node)
+
+    def parent(self, index):
+        if not index.isValid():
+            return QModelIndex()
+
+        child_node = index.internalPointer()
+        parent_node = child_node.parent
+
+        if parent_node == self.root:
+            return QModelIndex()
+
+        grand_parent = parent_node.parent
+        parent_row = grand_parent.children.index(parent_node)
+        return self.createIndex(parent_row, 0, parent_node)
+
+    def rowCount(self, index):
+        if index.column() > 0:
+            return 0
+
+        node = index.internalPointer() if index.isValid() else self.root
+        return len(node.children)
