Refactor the GetLegendGraphic file

lizmap_server/get_legend_graphic.py

@@ -10,7 +10,7 @@
 
 from typing import Optional
 
-from qgis.core import Qgis, QgsDataSourceUri, QgsProject
+from qgis.core import Qgis, QgsProject, QgsVectorLayer
 from qgis.server import QgsServerFilter
 
 from lizmap_server.core import find_vector_layer
@@ -19,8 +19,8 @@
 
 
 class GetLegendGraphicFilter(QgsServerFilter):
-    """add ruleKey to GetLegendGraphic for categorized and rule-based
-    only works for single LAYER and STYLE(S) and json format.
+    """ Add "ruleKey" to GetLegendGraphic for categorized and rule-based
+    only works for single LAYER and STYLE(S) and JSON format.
     """
 
     FEATURE_COUNT_REGEXP = r"(.*) \[≈?(?:\d+|N/A)\]"
@@ -87,53 +87,21 @@ def responseComplete(self):
                 if counter:
                     counter.waitForFinished()
 
-            renderer = layer.renderer()
-
             # From QGIS source code :
             # https://github.com/qgis/QGIS/blob/71499aacf431d3ac244c9b75c3d345bdc53572fb/src/core/symbology/qgsrendererregistry.cpp#L33
-            if renderer.type() in ("categorizedSymbol", "RuleRenderer", "graduatedSymbol"):
+            if layer.renderer().type() in ("categorizedSymbol", "RuleRenderer", "graduatedSymbol"):
                 body = handler.body()
                 # noinspection PyTypeChecker
                 json_data = json.loads(bytes(body))
-                categories = {}
-                for item in renderer.legendSymbolItems():
-
-                    # Calculate title if show_feature_count is activated
-                    # It seems that in QGIS Server 3.22 countSymbolFeatures is not used for JSON
-                    title = item.label()
-                    if show_feature_count:
-                        estimated_count = QgsDataSourceUri(layer.dataProvider().dataSourceUri()).useEstimatedMetadata()
-                        count = layer.featureCount(item.ruleKey())
-                        title += ' [{}{}]'.format(
-                            "≈" if estimated_count else "",
-                            count if count != -1 else "N/A",
-                        )
 
-                    expression = ''
-                    # TODO simplify when QGIS 3.26 will be the minimum version
-                    if Qgis.QGIS_VERSION_INT >= 32600:
-                        expression, result = renderer.legendKeyToExpression(item.ruleKey(), layer)
-                        if not result:
-                            Logger.warning(
-                                f"The expression in the project {project.homePath()}, layer {layer.name()} has not "
-                                f"been generated correctly, setting the expression to an empty string"
-                            )
-                            expression = ''
-
-                    categories[item.label()] = {
-                        'ruleKey': item.ruleKey(),
-                        'checked': renderer.legendSymbolItemChecked(item.ruleKey()),
-                        'parentRuleKey': item.parentRuleKey(),
-                        'scaleMaxDenom': item.scaleMaxDenom(),
-                        'scaleMinDenom': item.scaleMinDenom(),
-                        'expression': expression,
-                        'title': title,
-                    }
-
-                symbols = json_data['nodes'][0]['symbols'] if 'symbols' in json_data['nodes'][0] else json_data['nodes']
+                symbols = json_data['nodes'][0].get('symbols')
+                if not symbols:
+                    symbols = json_data['nodes']
 
                 new_symbols = []
 
+                categories = self._extract_categories(layer, show_feature_count, project.homePath())
+
                 for idx in range(len(symbols)):
                     symbol = symbols[idx]
                     symbol_label = symbol['title']
@@ -172,12 +140,58 @@ def responseComplete(self):
                     json_data['nodes'] = new_symbols
 
                 handler.clearBody()
-                handler.appendBody(json.dumps(
-                    json_data).encode('utf8'))
+                handler.appendBody(json.dumps(json_data).encode('utf8'))
         except Exception as ex:
             logger.critical(
                 'Error getting layer "{}" when setting up legend graphic for json output when configuring '
                 'OWS call: {}'.format(layer_name, str(ex)))
         finally:
             if layer and style and current_style and style != current_style:
                 layer.styleManager().setCurrentStyle(current_style)
+
+    @classmethod
+    def _extract_categories(
+            cls, layer: QgsVectorLayer, show_feature_count: bool = False, project_path: str = "") -> dict:
+        """ Extract categories from the layer legend. """
+        renderer = layer.renderer()
+        categories = {}
+        for item in renderer.legendSymbolItems():
+
+            # Calculate title if show_feature_count is activated
+            # It seems that in QGIS Server 3.22 countSymbolFeatures is not used for JSON
+            title = item.label()
+            if show_feature_count:
+                estimated_count = layer.dataProvider().uri().useEstimatedMetadata()
+                count = layer.featureCount(item.ruleKey())
+                title += ' [{}{}]'.format(
+                    "≈" if estimated_count else "",
+                    count if count != -1 else "N/A",
+                )
+
+            expression = ''
+            # TODO simplify when QGIS 3.26 will be the minimum version
+            if Qgis.QGIS_VERSION_INT >= 32600:
+                expression, result = renderer.legendKeyToExpression(item.ruleKey(), layer)
+                if not result:
+                    Logger.warning(
+                        f"The expression in the project {project_path}, layer {layer.name()} has not "
+                        f"been generated correctly, setting the expression to an empty string"
+                    )
+                    expression = ''
+
+            if item.label() in categories.keys():
+                Logger.warning(
+                    f"The label key '{item.label()}' is not unique, expect the legend to be broken in the project "
+                    f"{project_path}, layer {layer.name()}."
+                )
+
+            categories[item.label()] = {
+                'ruleKey': item.ruleKey(),
+                'checked': renderer.legendSymbolItemChecked(item.ruleKey()),
+                'parentRuleKey': item.parentRuleKey(),
+                'scaleMaxDenom': item.scaleMaxDenom(),
+                'scaleMinDenom': item.scaleMinDenom(),
+                'expression': expression,
+                'title': title,
+            }
+        return categories

test/test_legend.py

@@ -4,8 +4,6 @@
 
 from qgis.core import Qgis
 
-from lizmap_server.get_legend_graphic import GetLegendGraphicFilter
-
 LOGGER = logging.getLogger('server')
 
 __copyright__ = 'Copyright 2023, 3Liz'
@@ -139,15 +137,3 @@ def test_simple_rule_based_feature_count(client):
     assert symbols[0]['expression'] == expected, symbols[0]['expression']
     assert b['title'] == ''
     assert b['nodes'][0]['title'] == 'rule_based [4]', b['nodes'][0]['title']
-
-
-def test_regexp_feature_count():
-    """ Test the regexp about the feature count. """
-    result = GetLegendGraphicFilter.match_label_feature_count("A label [22]")
-    assert result.group(1) == "A label", result.group(1)
-
-    result = GetLegendGraphicFilter.match_label_feature_count("A label [≈2]")
-    assert result.group(1) == "A label", result.group(1)
-
-    result = GetLegendGraphicFilter.match_label_feature_count("A label")
-    assert result is None

test/test_legend_without_server.py

@@ -0,0 +1,75 @@
+import unittest
+
+from qgis.core import (
+    Qgis,
+    QgsRuleBasedRenderer,
+    QgsSymbol,
+    QgsVectorLayer,
+    QgsWkbTypes,
+)
+
+from lizmap_server.get_legend_graphic import GetLegendGraphicFilter
+
+__copyright__ = 'Copyright 2024, 3Liz'
+__license__ = 'GPL version 3'
+__email__ = '[email protected]'
+
+
+class TestLegend(unittest.TestCase):
+
+    def test_regexp_feature_count(self):
+        """ Test the regexp about the feature count. """
+        result = GetLegendGraphicFilter.match_label_feature_count("A label [22]")
+        self.assertEqual(result.group(1), "A label")
+
+        result = GetLegendGraphicFilter.match_label_feature_count("A label [≈2]")
+        self.assertEqual(result.group(1), "A label")
+
+        result = GetLegendGraphicFilter.match_label_feature_count("A label")
+        self.assertIsNone(result)
+
+    def test_duplicated_labels(self):
+        """ Test the legend with multiple sub-rules in the rule based rendered. """
+        # noinspection PyTypeChecker
+        root_rule = QgsRuleBasedRenderer.Rule(None)
+
+        same_label = 'same-label'
+
+        # Rule 1 with symbol
+        # noinspection PyUnresolvedReferences
+        rule_1 = QgsRuleBasedRenderer.Rule(QgsSymbol.defaultSymbol(QgsWkbTypes.PointGeometry), label='rule-1')
+        root_rule.appendChild(rule_1)
+
+        # Sub-rule to rule 1
+        # noinspection PyTypeChecker
+        rule_1_1 = QgsRuleBasedRenderer.Rule(None, label=same_label)
+        rule_1.appendChild(rule_1_1)
+
+        # Rule 2 with symbol
+        # noinspection PyUnresolvedReferences
+        rule_2 = QgsRuleBasedRenderer.Rule(QgsSymbol.defaultSymbol(QgsWkbTypes.PointGeometry), label='rule-2')
+        root_rule.appendChild(rule_2)
+
+        # Sub-rule to rule 2
+        # noinspection PyTypeChecker
+        rule_2_1 = QgsRuleBasedRenderer.Rule(None, label=same_label)
+        rule_2.appendChild(rule_2_1)
+
+        layer = QgsVectorLayer("Point?field=fldtxt:string", "layer1", "memory")
+        layer.setRenderer(QgsRuleBasedRenderer(root_rule))
+
+        result = GetLegendGraphicFilter._extract_categories(layer)
+        # TODO, this should be 4, as we have 4 rules
+        self.assertEqual(3, len(list(result.keys())))
+
+        for symbol in result.values():
+            self.assertGreaterEqual(len(symbol['ruleKey']), 1)
+            self.assertTrue(symbol['checked'])
+            self.assertGreaterEqual(len(symbol['parentRuleKey']), 1)
+            self.assertEqual(0, symbol['scaleMaxDenom'])
+            self.assertEqual(0, symbol['scaleMinDenom'])
+            if Qgis.QGIS_VERSION_INT >= 33400:
+                self.assertEqual('TRUE', symbol['expression'])
+            else:
+                self.assertEqual('', symbol['expression'])
+            self.assertIn(symbol['title'], ('rule-1', 'same-label', 'rule-2'))