fix: fix type error in assignment with type annotation (kcl-lang#806)

Signed-off-by: zongz <[email protected]>
Peefy · Oct 26, 2023 · 9d36d96 · 9d36d96
1 parent 1d9eccb
commit 9d36d96
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 4 deletions.
diff --git a/kclvm/sema/src/resolver/node.rs b/kclvm/sema/src/resolver/node.rs
@@ -160,11 +160,16 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Resolver<'ctx> {
                     }
                     TypeKind::Dict(_) => {
                         value_ty = self.expr(&assign_stmt.value);
-                        // update attrs
+                        if !assign_stmt.type_annotation.is_none() {
+                            // Check type annotation if exists.
+                            self.check_assignment_type_annotation(assign_stmt, value_ty.clone());
+                        }
                         self.set_type_to_scope(name, value_ty.clone(), target.get_span_pos());
                     }
                     _ => {
                         value_ty = self.expr(&assign_stmt.value);
+                        // Check type annotation if exists.
+                        self.check_assignment_type_annotation(assign_stmt, value_ty.clone());
                     }
                 }
                 self.must_assignable_to(
@@ -192,12 +197,11 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Resolver<'ctx> {
                 let expected_ty = self.walk_identifier_expr(target);
                 self.ctx.l_value = false;
                 value_ty = self.expr(&assign_stmt.value);
+                // Check type annotation if exists.
+                self.check_assignment_type_annotation(assign_stmt, value_ty.clone());
                 self.must_assignable_to(value_ty.clone(), expected_ty, target.get_span_pos(), None)
             }
         }
-        // Check type annotation if exists.
-        self.check_assignment_type_annotation(assign_stmt, value_ty.clone());
-
         value_ty
     }
 

diff --git a/kclvm/sema/src/resolver/test_data/annotation_check_assignment.k b/kclvm/sema/src/resolver/test_data/annotation_check_assignment.k
@@ -0,0 +1,6 @@
+params = option("params") or {}
+# Use `k = v` to override existing annotations
+annotations: {str:str} = {k = v for k, v in params.annotations or {}}
+items = [item | {
+    metadata.annotations: annotations
+} for item in option("items")]
diff --git a/kclvm/sema/src/resolver/tests.rs b/kclvm/sema/src/resolver/tests.rs
@@ -6,8 +6,10 @@ use crate::resolver::resolve_program;
 use crate::resolver::scope::*;
 use crate::ty::{Type, TypeKind};
 use kclvm_ast::ast;
+use kclvm_ast::ast::CmdArgSpec;
 use kclvm_ast::pos::ContainsPos;
 use kclvm_error::*;
+use kclvm_parser::LoadProgramOptions;
 use kclvm_parser::ParseSession;
 use kclvm_parser::{load_program, parse_program};
 use std::path::Path;
@@ -550,6 +552,30 @@ fn test_resolve_assignment_in_lambda() {
     assert_eq!(images_scope_obj.borrow().ty.ty_str(), "[str]");
 }
 
+#[test]
+fn test_assignment_type_annotation_check_in_lambda() {
+    let sess = Arc::new(ParseSession::default());
+    let mut opts = LoadProgramOptions::default();
+    opts.cmd_args = vec![
+        CmdArgSpec {
+            name: "params".to_string(),
+            value: "annotations: {a: b}".to_string(),
+        },
+        CmdArgSpec {
+            name: "items".to_string(),
+            value: "metadata: {annotations:{c: d}}".to_string(),
+        },
+    ];
+    let mut program = load_program(
+        sess.clone(),
+        &["./src/resolver/test_data/annotation_check_assignment.k"],
+        Some(opts),
+    )
+    .unwrap();
+    let scope = resolve_program(&mut program);
+    assert_eq!(scope.handler.diagnostics.len(), 0);
+}
+
 #[test]
 fn test_resolve_lambda_assignment_diagnostic() {
     let sess = Arc::new(ParseSession::default());