From deb6736f9fc5267e65bb4b8a5f8217330aa9156c Mon Sep 17 00:00:00 2001
From: Shunsuke Shibayama <sbym1346@gmail.com>
Date: Fri, 27 Dec 2024 20:09:50 +0900
Subject: [PATCH] feat: support named-expr (walrus operator)

---
 crates/py2erg/convert.rs | 45 ++++++++++++++++++++++++++++++----------
 tests/test.py            |  3 +++
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/crates/py2erg/convert.rs b/crates/py2erg/convert.rs
index 0d2c16c..f0b7517 100644
--- a/crates/py2erg/convert.rs
+++ b/crates/py2erg/convert.rs
@@ -10,17 +10,17 @@ use erg_common::traits::{Locational, Stream};
 use erg_common::{fmt_vec, log, set};
 use erg_compiler::artifact::IncompleteArtifact;
 use erg_compiler::erg_parser::ast::{
-    Accessor, Args, BinOp, Block, ClassAttr, ClassAttrs, ClassDef, ConstAccessor, ConstApp,
-    ConstArgs, ConstAttribute, ConstBinOp, ConstBlock, ConstDict, ConstExpr, ConstKeyValue,
-    ConstLambda, ConstList, ConstListWithLength, ConstNormalList, ConstNormalSet, ConstPosArg,
-    ConstSet, Decorator, Def, DefBody, DefId, DefaultParamSignature, Dict, Dummy, Expr, Identifier,
-    KeyValue, KwArg, Lambda, LambdaSignature, List, ListComprehension, Literal, Methods, Module,
-    NonDefaultParamSignature, NormalDict, NormalList, NormalRecord, NormalSet, NormalTuple,
-    ParamPattern, ParamTySpec, Params, PosArg, PreDeclTypeSpec, ReDef, Record, RecordAttrs, Set,
-    SetComprehension, Signature, SubrSignature, SubrTypeSpec, Tuple, TupleTypeSpec, TypeAppArgs,
-    TypeAppArgsKind, TypeAscription, TypeBoundSpec, TypeBoundSpecs, TypeSpec, TypeSpecWithOp,
-    UnaryOp, VarName, VarPattern, VarRecordAttr, VarRecordAttrs, VarRecordPattern, VarSignature,
-    VisModifierSpec,
+    Accessor, Args, BinOp, Block, ClassAttr, ClassAttrs, ClassDef, Compound, ConstAccessor,
+    ConstApp, ConstArgs, ConstAttribute, ConstBinOp, ConstBlock, ConstDict, ConstExpr,
+    ConstKeyValue, ConstLambda, ConstList, ConstListWithLength, ConstNormalList, ConstNormalSet,
+    ConstPosArg, ConstSet, Decorator, Def, DefBody, DefId, DefaultParamSignature, Dict, Dummy,
+    Expr, Identifier, KeyValue, KwArg, Lambda, LambdaSignature, List, ListComprehension, Literal,
+    Methods, Module, NonDefaultParamSignature, NormalDict, NormalList, NormalRecord, NormalSet,
+    NormalTuple, ParamPattern, ParamTySpec, Params, PosArg, PreDeclTypeSpec, ReDef, Record,
+    RecordAttrs, Set, SetComprehension, Signature, SubrSignature, SubrTypeSpec, Tuple,
+    TupleTypeSpec, TypeAppArgs, TypeAppArgsKind, TypeAscription, TypeBoundSpec, TypeBoundSpecs,
+    TypeSpec, TypeSpecWithOp, UnaryOp, VarName, VarPattern, VarRecordAttr, VarRecordAttrs,
+    VarRecordPattern, VarSignature, VisModifierSpec,
 };
 use erg_compiler::erg_parser::desugar::Desugarer;
 use erg_compiler::erg_parser::token::{Token, TokenKind, AS, COLON, DOT, EQUAL};
@@ -1948,6 +1948,29 @@ impl ASTConverter {
                 let stringify = self.convert_ident("str".to_string(), loc);
                 stringify.call1(expr).into()
             }
+            py_ast::Expr::NamedExpr(named) => {
+                let loc = named.location();
+                let target = self.convert_expr(*named.target);
+                let target_pat = match &target {
+                    Expr::Accessor(Accessor::Ident(ident)) => VarPattern::Ident(ident.clone()),
+                    _ => {
+                        log!(err "unimplemented: {:?}", target);
+                        VarPattern::Ident(Identifier::private("_".into()))
+                    }
+                };
+                let value = self.convert_expr(*named.value);
+                let assign = Token::new(
+                    TokenKind::Assign,
+                    "=",
+                    loc.row.get(),
+                    loc.column.to_zero_indexed(),
+                );
+                let def = Def::new(
+                    Signature::Var(VarSignature::new(target_pat, None)),
+                    DefBody::new(assign, Block::new(vec![value]), DefId(0)),
+                );
+                Expr::Compound(Compound::new(vec![Expr::Def(def), target]))
+            }
             _other => {
                 log!(err "unimplemented: {:?}", _other);
                 Expr::Dummy(Dummy::new(None, vec![]))
diff --git a/tests/test.py b/tests/test.py
index 9b2c359..4e234ad 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -44,3 +44,6 @@ assert j == 2
 with open("test.py") as f:
     for line in f.readlines():
         print("line: " + line)
+
+print(x := 1)
+print(x)