From c564153192dc4c7abb501856c8f82226a5da0be3 Mon Sep 17 00:00:00 2001
From: "Leslie H." <che27@illinois.edu>
Date: Wed, 12 Jun 2024 15:33:14 -0500
Subject: [PATCH] Update dropdown conditions on column rename

---
 sandbox/grist/dropdown_condition.py | 74 ++++++++++++++++++++++++++++-
 sandbox/grist/useractions.py        |  3 ++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/sandbox/grist/dropdown_condition.py b/sandbox/grist/dropdown_condition.py
index b57ef3f0542..648609e0802 100644
--- a/sandbox/grist/dropdown_condition.py
+++ b/sandbox/grist/dropdown_condition.py
@@ -1,10 +1,55 @@
+import ast
+import asttokens
 import json
 import logging
+import textbuilder
 
-from predicate_formula import parse_predicate_formula_json
+from predicate_formula import NamedEntity, parse_predicate_formula_json, TreeConverter
 
 log = logging.getLogger(__name__)
 
+
+def perform_dropdown_condition_renames(useractions, renames):
+  """
+  Given a dict of column renames of the form {(table_id, col_id): new_col_id}, applys updates
+  to the affected dropdown condition formulas.
+  """
+  updates = []
+
+  for col in useractions._engine.docmodel.columns.all:
+
+    patches = []
+
+    # Find all columns in the document that has dropdown conditions.
+    try:
+      widget_options = json.loads(col.widgetOptions)
+      dropdown_condition = widget_options["dropdownCondition"]["text"]
+    except Exception as e:
+      continue
+    # Find out what table this column refers to.
+    table_id = col.type.lstrip("Ref:")
+
+    # Parse the formula to find out all attributes that could be subject to rename.
+    for subject in parse_dropdown_condition_grist_entities(dropdown_condition):
+      # Only process attributes that represent a renamed column.
+      if (table_id, subject.name) in renames:
+        old_name = subject.name
+        new_name = renames[(table_id, subject.name)]
+        patches.append(textbuilder.make_patch(
+            dropdown_condition, subject.start_pos, subject.start_pos + len(old_name), new_name))
+
+    # Replace the column reference with the new name.
+    new_dropdown_condition = textbuilder.Replacer(textbuilder.Text(dropdown_condition), patches).get_text()
+
+    # Parse the new dropdown condition formula.
+    widget_options["dropdownCondition"] = {"text": new_dropdown_condition,
+                                           "parsed": parse_predicate_formula_json(new_dropdown_condition)}
+    updates.append((col, {"widgetOptions": json.dumps(widget_options)}))
+
+    # Update the dropdown condition in the database.
+    useractions.doBulkUpdateFromPairs('_grist_Tables_column', updates)
+
+
 def parse_dropdown_conditions(col_values):
   """
   Parses any unparsed dropdown conditions in `col_values`.
@@ -41,3 +86,30 @@ def parse_dropdown_condition(widget_options_json):
     return json.dumps(widget_options)
   except (TypeError, ValueError):
     return widget_options_json
+
+
+def parse_dropdown_condition_grist_entities(dc_formula):
+  """
+  Parse the dropdown condition formula collecting any entities that may be subject to renaming.
+  Returns a NamedEntity list.
+  See also: parse_acl_grist_entities
+  """
+  try:
+    atok = asttokens.ASTTokens(dc_formula, tree=ast.parse(dc_formula, mode='eval'))
+    converter = _EntityCollector()
+    converter.visit(atok.tree)
+    return converter.entities
+  except SyntaxError as err:
+    return []
+
+class _EntityCollector(TreeConverter):
+  def __init__(self):
+    self.entities = []
+
+  def visit_Attribute(self, node):
+    parent = self.visit(node.value)
+
+    if parent == ["Name", "choice"]:
+      self.entities.append(NamedEntity(None, node.last_token.startpos, node.attr, None))
+
+    return ["Attr", parent, node.attr]
diff --git a/sandbox/grist/useractions.py b/sandbox/grist/useractions.py
index ac9261ef1a1..36c7641efcc 100644
--- a/sandbox/grist/useractions.py
+++ b/sandbox/grist/useractions.py
@@ -14,6 +14,7 @@
 import gencode
 from acl_formula import parse_acl_formulas
 from dropdown_condition import parse_dropdown_conditions
+import dropdown_condition
 import actions
 import column
 import sort_specs
@@ -720,6 +721,8 @@ def _updateColumnRecords(self, table_id, row_ids, col_values):
 
     make_acl_updates = acl.prepare_acl_col_renames(self._docmodel, self, renames)
 
+    dropdown_condition.perform_dropdown_condition_renames(self, renames)
+
     rename_summary_tables = set()
     for c, values in update_pairs:
       # Trigger ModifyColumn and RenameColumn as necessary