diff --git a/sandbox/grist/codebuilder.py b/sandbox/grist/codebuilder.py index 43b736b1..2309333d 100644 --- a/sandbox/grist/codebuilder.py +++ b/sandbox/grist/codebuilder.py @@ -247,6 +247,23 @@ class InferLookupReference(InferenceTip): def infer(cls, node, context=None): yield astroid.bases.Instance(infer(node.func.expr)) + +class InferAllReference(InferenceTip): + """ + Inference helper to treat the return value of `Table.all` as returning instances + of table `Table`. + """ + node_class = astroid.nodes.Attribute + + @classmethod + def filter(cls, node): + return node.attrname == "all" and _is_table(infer(node.expr)) + + @classmethod + def infer(cls, node, context=None): + yield astroid.bases.Instance(infer(node.expr)) + + class InferLookupComprehension(InferenceTip): node_class = astroid.nodes.AssignName @@ -255,9 +272,11 @@ class InferLookupComprehension(InferenceTip): compr = node.parent if not isinstance(compr, astroid.nodes.Comprehension): return False - if not isinstance(compr.iter, astroid.nodes.Call): - return False - return InferLookupReference.filter(compr.iter) + if isinstance(compr.iter, astroid.nodes.Call): + return InferLookupReference.filter(compr.iter) + if isinstance(compr.iter, astroid.nodes.Attribute): + return InferAllReference.filter(compr.iter) + return False @classmethod def infer(cls, node, context=None): @@ -311,7 +330,7 @@ def parse_grist_names(builder): code_text = builder.get_text() with use_inferences(InferReferenceColumn, InferReferenceFormula, InferLookupReference, - InferLookupComprehension): + InferLookupComprehension, InferAllReference): 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 de8b92bd..5ff6085a 100644 --- a/sandbox/grist/test_renames.py +++ b/sandbox/grist/test_renames.py @@ -213,6 +213,7 @@ class TestRenames(test_engine.EngineTestCase): self.load_sample(self.sample) self.add_column("Address", "people", formula="People.lookupOne(addr=$id, city=$city).name") self.add_column("Address", "people2", formula="People.lookupRecords(addr=$id).name") + self.add_column("Address", "people3", formula="People.all.name") out_actions = self.apply_user_action(["RenameColumn", "People", "name", "nombre"]) self.assertPartialOutActions(out_actions, { "stored": [ ["RenameColumn", "People", "name", "nombre"], @@ -220,11 +221,14 @@ class TestRenames(test_engine.EngineTestCase): "People.lookupOne(addr=$id, city=$city).nombre"}], ["ModifyColumn", "Address", "people2", {"formula": "People.lookupRecords(addr=$id).nombre"}], - ["BulkUpdateRecord", "_grist_Tables_column", [22, 25, 26], { - "colId": ["nombre", "people", "people2"], + ["ModifyColumn", "Address", "people3", {"formula": + "People.all.nombre"}], + ["BulkUpdateRecord", "_grist_Tables_column", [22, 25, 26, 27], { + "colId": ["nombre", "people", "people2", "people3"], "formula": ["", "People.lookupOne(addr=$id, city=$city).nombre", - "People.lookupRecords(addr=$id).nombre"] + "People.lookupRecords(addr=$id).nombre", + "People.all.nombre"] }], ]})