From 525613216c234bf70b4835a98b5741938715c6ca Mon Sep 17 00:00:00 2001
From: Alex Hall <alex.mojaki@gmail.com>
Date: Thu, 7 Sep 2023 12:01:42 +0200
Subject: [PATCH] (core) Fix updating attributes inside f-strings when columns
 are renamed

Summary: Upgrades asttokens and uses the newly released support for astroid trees in `asttokens.ASTText` which is needed to deal with f-strings (as opposed to `asttokens.ASTTokens`).

Test Plan: Added a Python unit test

Reviewers: dsagal

Reviewed By: dsagal

Differential Revision: https://phab.getgrist.com/D4027
---
 sandbox/grist/codebuilder.py  | 18 ++++++++++--------
 sandbox/grist/test_renames.py | 17 +++++++++++++++++
 sandbox/requirements.txt      |  2 +-
 sandbox/requirements3.txt     |  2 +-
 4 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/sandbox/grist/codebuilder.py b/sandbox/grist/codebuilder.py
index 66f7cd3e..8bf4746f 100644
--- a/sandbox/grist/codebuilder.py
+++ b/sandbox/grist/codebuilder.py
@@ -397,7 +397,7 @@ def parse_grist_names(builder):
 
   with use_inferences(InferReferenceColumn, InferReferenceFormula, InferLookupReference,
                       InferLookupComprehension, InferAllReference, InferAllComprehension):
-    atok = asttokens.ASTTokens(code_text, tree=astroid.builder.parse(code_text))
+    atok = asttokens.ASTText(code_text, tree=astroid.builder.parse(code_text))
 
   def make_tuple(start, end, table_id, col_id):
     name = col_id or table_id
@@ -413,7 +413,7 @@ def parse_grist_names(builder):
     return None
 
   parsed_names = []
-  for node in asttokens.util.walk(atok.tree):
+  for node in asttokens.util.walk(atok.tree, include_joined_str=True):
     if isinstance(node, astroid.nodes.Name):
       obj = infer(node)
       if _is_table(obj) and not _is_local(node):
@@ -425,17 +425,19 @@ def parse_grist_names(builder):
       if isinstance(obj, astroid.bases.Instance):
         cls = obj._proxied
         if _is_table(cls):
-          tok = node.last_token
-          start, end = tok.startpos, tok.endpos
-          parsed_names.append(make_tuple(start, end, cls.name, node.attrname))
+          end = atok.get_text_range(node)[1]
+          start = end - len(node.attrname)
+          if code_text[start:end] == node.attrname:
+            parsed_names.append(make_tuple(start, end, cls.name, node.attrname))
     elif isinstance(node, astroid.nodes.Keyword):
       func = node.parent.func
       if isinstance(func, astroid.nodes.Attribute) and func.attrname in _lookup_method_names:
         obj = infer(func.expr)
         if _is_table(obj):
-          tok = node.first_token
-          start, end = tok.startpos, tok.endpos
-          parsed_names.append(make_tuple(start, end, obj.name, node.arg))
+          start = atok.get_text_range(node)[0]
+          end = start + len(node.arg)
+          if code_text[start:end] == node.arg:
+            parsed_names.append(make_tuple(start, end, obj.name, node.arg))
 
   return [name for name in parsed_names if name]
 
diff --git a/sandbox/grist/test_renames.py b/sandbox/grist/test_renames.py
index 5fac8f31..b7ae31e5 100644
--- a/sandbox/grist/test_renames.py
+++ b/sandbox/grist/test_renames.py
@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
 import logging
+import unittest
+
+import six
 
 import testutil
 import test_engine
@@ -78,6 +81,20 @@ class TestRenames(test_engine.EngineTestCase):
       }]
     ]})
 
+  @unittest.skipUnless(six.PY3, "Python 3 only")
+  def test_rename_inside_fstring(self):
+    self.load_sample(self.sample)
+    self.add_column("People", "CityUpper", formula="f'{$city.upper()}'")
+    out_actions = self.apply_user_action(["RenameColumn", "People", "city", "ciudad"])
+    self.assertPartialOutActions(out_actions, { "stored": [
+      ["RenameColumn", "People", "city", "ciudad"],
+      ["ModifyColumn", "People", "CityUpper", {"formula": "f'{$ciudad.upper()}'"}],
+      ["BulkUpdateRecord", "_grist_Tables_column", [24, 25], {
+        "colId": ["ciudad", "CityUpper"],
+        "formula": ["$addr.city", "f'{$ciudad.upper()}'"]
+      }]
+    ]})
+
   def test_rename_reference_attribute(self):
     # Slightly harder: renaming `$ref.COLUMN`
     self.load_sample(self.sample)
diff --git a/sandbox/requirements.txt b/sandbox/requirements.txt
index 118531ac..40ad6c77 100644
--- a/sandbox/requirements.txt
+++ b/sandbox/requirements.txt
@@ -1,7 +1,7 @@
 ### python 2 requirements, see requirements3.txt for python 3
 
 astroid==1.6.6
-asttokens==2.2.1
+asttokens==2.4.0
 backports.functools-lru-cache==1.6.4
 chardet==4.0.0
 enum34==1.1.10
diff --git a/sandbox/requirements3.txt b/sandbox/requirements3.txt
index d8168ae3..b7f84e42 100644
--- a/sandbox/requirements3.txt
+++ b/sandbox/requirements3.txt
@@ -6,7 +6,7 @@
 #
 astroid==2.14.2
     # via -r core/sandbox/requirements3.in
-asttokens==2.2.1
+asttokens==2.4.0
     # via
     #   friendly-traceback
     #   stack-data