(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
This commit is contained in:
Alex Hall 2022-07-21 18:02:01 +02:00
parent 460f22b701
commit 938928f1b9
2 changed files with 35 additions and 19 deletions

View File

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

View File

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