LineaLabs/lineapy

Support walrus operator `:=`

yifanwu opened this issue · 1 comments

What: Walrus operator, e.g., x = (y := 0), currently breaks our node_transformer assumptions.

Why: the shape of the AST looks different:

~/linea-dev/lineapy/lineapy/ipython.py in _end_cell()
    146     get_cell_path(location).write_text(code)
    147 
--> 148     last_node = transform(code, location, STATE.tracer)
    149     if STATE.visualize_display_handle:
    150         STATE.visualize_display_handle.update(

~/linea-dev/lineapy/lineapy/transformer/node_transformer.py in transform(code, location, tracer)
     80     except SyntaxError as e:
     81         raise UserException(e, RemoveFrames(2))
---> 82     node_transformer.visit(tree)
     83 
     84     tracer.db.commit()

~/anaconda3/envs/lineapy-env/lib/python3.9/ast.py in visit(self, node)
    405         method = 'visit_' + node.__class__.__name__
    406         visitor = getattr(self, method, self.generic_visit)
--> 407         return visitor(node)
    408 
    409     def generic_visit(self, node):

~/linea-dev/lineapy/lineapy/transformer/node_transformer.py in visit_Module(self, node)
    153     def visit_Module(self, node: ast.Module) -> Any:
    154         for stmt in node.body:
--> 155             self.last_statement_result = self.visit(stmt)
    156 
    157     def visit_Expr(self, node: ast.Expr) -> Node:

~/anaconda3/envs/lineapy-env/lib/python3.9/ast.py in visit(self, node)
    405         method = 'visit_' + node.__class__.__name__
    406         visitor = getattr(self, method, self.generic_visit)
--> 407         return visitor(node)
    408 
    409     def generic_visit(self, node):

~/linea-dev/lineapy/lineapy/transformer/node_transformer.py in visit_Assign(self, node)
    261         self.visit_assign_value(
    262             target,
--> 263             self.visit(node.value),
    264             self.get_source(node),
    265         )

~/anaconda3/envs/lineapy-env/lib/python3.9/ast.py in visit(self, node)
    405         method = 'visit_' + node.__class__.__name__
    406         visitor = getattr(self, method, self.generic_visit)
--> 407         return visitor(node)
    408 
    409     def generic_visit(self, node):

~/linea-dev/lineapy/lineapy/transformer/node_transformer.py in generic_visit(self, node)
    134             return self._exec_expression(node)
    135         else:
--> 136             raise NotImplementedError(
    137                 f"Don't know how to transform {type(node).__name__}"
    138             )

NotImplementedError: Don't know how to transform NamedExpr

Possible Approaches: this falls in the gray area of eval/syntax. Might be similar to the star operator (where we just create a short-cut in node_transformer), or we might have to make a custom node and pass to the executor. Anyhow this should be a lower priority item.

This has been done!