From f5eaff35c50efbc16b784fdf3d50f458cf5a2ef3 Mon Sep 17 00:00:00 2001
From: Maxence Laurent <6583528+maxencelaurent@users.noreply.github.com>
Date: Sun, 23 Jun 2024 16:24:04 +0200
Subject: [PATCH] fix optimize_expression

* allow non-numeric PK expression optimization
* do not assume PK is the very first field of the list: use
  layer.primaryKeyAttributes() instead
* fixes 3liz/lizmap-web-client#4534
* related to 3liz/qgis-atlasprint#71
---
 atlasprint/core.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/atlasprint/core.py b/atlasprint/core.py
index f782186..4dedea4 100755
--- a/atlasprint/core.py
+++ b/atlasprint/core.py
@@ -312,17 +312,19 @@ def optimize_expression(layer, expression):
         logger.info("'$id' not found in the expression, returning the input expression.")
         return expression
 
+    # extract indexes of primary keys
     primary_keys = layer.primaryKeyAttributes()
     if len(primary_keys) != 1:
         logger.info("Primary keys are not defined in the layer '{}'.".format(layer.id()))
         return expression
 
-    field = layer.fields().at(0)
-    if not field.isNumeric():
-        logger.info("The field '{}' is not numeric in layer '{}'.".format(field.name(), layer.id()))
-        return expression
 
-    expression = expression.replace('$id', '"{}"'.format(field.name()))
-    logger.info('$id has been replaced by "{}" in layer "{}"'.format(field.name(), layer.id()))
+    # extract primary key from fields list
+    pk_index = primary_keys[0];
+    pk_field = layer.fields().at(pk_index)
+
+    # replace `$id` with effective PK name
+    expression = expression.replace('$id', '"{}"'.format(pk_field.name()))
+    logger.info('$id has been replaced by "{}" in layer "{}"'.format(pk_field.name(), layer.id()))
 
     return expression