(core) Fix renaming columns when iterating over Table.all in comprehension

Summary:
Comprehensions iterating over `Table.all` like `[foo.bar for foo in Table.all]` led to an error when renaming the column `bar`. This diff fixes that so that renaming `bar` does the same thing as for a comprehension over `Table.lookupRecords()`. Note that `next(foo for foo in Table.all).bar` is still not supported, as the same is not supported for `Table.lookupRecords()` either.

Discussion: https://grist.slack.com/archives/C069RUP71/p1658360276762949

Test Plan: Parametrised existing Python test to test the same thing for both `all` and `lookupRecords`

Reviewers: dsagal

Reviewed By: dsagal

Subscribers: dsagal

Differential Revision: https://phab.getgrist.com/D3538

sandbox/grist/codebuilder.py

@@ -264,23 +264,31 @@ class InferAllReference(InferenceTip):
     yield astroid.bases.Instance(infer(node.expr))
 
 
-class InferLookupComprehension(InferenceTip):
+class InferComprehensionBase(InferenceTip):
   node_class = astroid.nodes.AssignName
+  reference_inference_class = None
 
   @classmethod
   def filter(cls, node):
     compr = node.parent
     if not isinstance(compr, astroid.nodes.Comprehension):
       return False
-    if isinstance(compr.iter, astroid.nodes.Call):
-      return InferLookupReference.filter(compr.iter)
-    if isinstance(compr.iter, astroid.nodes.Attribute):
-      return InferAllReference.filter(compr.iter)
+    if isinstance(compr.iter, cls.reference_inference_class.node_class):
+      return cls.reference_inference_class.filter(compr.iter)
     return False
 
   @classmethod
   def infer(cls, node, context=None):
-    return InferLookupReference.infer(node.parent.iter)
+    return cls.reference_inference_class.infer(node.parent.iter)
+
+
+class InferLookupComprehension(InferComprehensionBase):
+  reference_inference_class = InferLookupReference
+
+
+class InferAllComprehension(InferComprehensionBase):
+  reference_inference_class = InferAllReference
+
 
 class InferRecAssignment(InferenceTip):
   """
@@ -330,7 +338,7 @@ def parse_grist_names(builder):
   code_text = builder.get_text()
 
   with use_inferences(InferReferenceColumn, InferReferenceFormula, InferLookupReference,
-                      InferLookupComprehension, InferAllReference):
+                      InferLookupComprehension, InferAllReference, InferAllComprehension):
     atok = asttokens.ASTTokens(code_text, tree=astroid.builder.parse(code_text))
 
   def make_tuple(start, end, table_id, col_id):

sandbox/grist/test_renames.py

@@ -255,34 +255,42 @@ class TestRenames(test_engine.EngineTestCase):
 
   def test_rename_lookup_iter_attr(self):
     # Renaming `[x.COLUMN for x in Table.lookupRecords(...)]`.
+    self.check_comprehension_rename("People.lookupRecords(addr=$id)",
+                                    "People.lookupRecords(ADDRESS=$id)")
+
+  def test_rename_all_iter_attr(self):
+    # Renaming `[x.COLUMN for x in Table.all]`.
+    self.check_comprehension_rename("People.all", "People.all")
+
+  def check_comprehension_rename(self, iter_expr1, iter_expr2):
     self.load_sample(self.sample)
     self.add_column("Address", "people",
-                    formula="','.join(x.addr.city for x in People.lookupRecords(addr=$id))")
+                    formula="','.join(x.addr.city for x in %s)" % iter_expr1)
     self.add_column("Address", "people2",
-                    formula="','.join([x.addr.city for x in People.lookupRecords(addr=$id)])")
+                    formula="','.join([x.addr.city for x in %s])" % iter_expr1)
     self.add_column("Address", "people3",
-                    formula="','.join({x.addr.city for x in People.lookupRecords(addr=$id)})")
+                    formula="','.join({x.addr.city for x in %s})" % iter_expr1)
     self.add_column("Address", "people4",
-                    formula="{x.addr.city:x.addr for x in People.lookupRecords(addr=$id)}")
+                    formula="{x.addr.city:x.addr for x in %s}" % iter_expr1)
     out_actions = self.apply_user_action(["RenameColumn", "People", "addr", "ADDRESS"])
     self.assertPartialOutActions(out_actions, { "stored": [
       ["RenameColumn", "People", "addr", "ADDRESS"],
       ["ModifyColumn", "People", "city", {"formula": "$ADDRESS.city"}],
       ["ModifyColumn", "Address", "people",
-           {"formula": "','.join(x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id))"}],
+           {"formula": "','.join(x.ADDRESS.city for x in %s)" % iter_expr2}],
       ["ModifyColumn", "Address", "people2",
-           {"formula": "','.join([x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id)])"}],
+           {"formula": "','.join([x.ADDRESS.city for x in %s])" % iter_expr2}],
       ["ModifyColumn", "Address", "people3",
-           {"formula": "','.join({x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id)})"}],
+           {"formula": "','.join({x.ADDRESS.city for x in %s})" % iter_expr2}],
       ["ModifyColumn", "Address", "people4",
-           {"formula": "{x.ADDRESS.city:x.ADDRESS for x in People.lookupRecords(ADDRESS=$id)}"}],
+           {"formula": "{x.ADDRESS.city:x.ADDRESS for x in %s}" % iter_expr2}],
       ["BulkUpdateRecord", "_grist_Tables_column", [23, 24, 25, 26, 27, 28], {
         "colId": ["ADDRESS", "city", "people", "people2", "people3", "people4"],
         "formula": ["", "$ADDRESS.city",
-           "','.join(x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id))",
-           "','.join([x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id)])",
-           "','.join({x.ADDRESS.city for x in People.lookupRecords(ADDRESS=$id)})",
-           "{x.ADDRESS.city:x.ADDRESS for x in People.lookupRecords(ADDRESS=$id)}"],
+           "','.join(x.ADDRESS.city for x in %s)" % iter_expr2,
+           "','.join([x.ADDRESS.city for x in %s])" % iter_expr2,
+           "','.join({x.ADDRESS.city for x in %s})" % iter_expr2,
+           "{x.ADDRESS.city:x.ADDRESS for x in %s}" % iter_expr2],
       }],
     ]})