From 938928f1b907f4364ce5dcd0702dbfeda33cdc3e Mon Sep 17 00:00:00 2001 From: Alex Hall Date: Thu, 21 Jul 2022 18:02:01 +0200 Subject: [PATCH] (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 | 22 +++++++++++++++------- sandbox/grist/test_renames.py | 32 ++++++++++++++++++++------------ 2 files changed, 35 insertions(+), 19 deletions(-) diff --git a/sandbox/grist/codebuilder.py b/sandbox/grist/codebuilder.py index 2309333d..3466a278 100644 --- a/sandbox/grist/codebuilder.py +++ b/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): diff --git a/sandbox/grist/test_renames.py b/sandbox/grist/test_renames.py index 5ff6085a..4e4f91c6 100644 --- a/sandbox/grist/test_renames.py +++ b/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], }], ]})