feat: impl subscript in override spec (#1026)

Signed-off-by: peefy <[email protected]>

kclvm/ast_pretty/src/node.rs

@@ -847,7 +847,7 @@ impl<'p> Printer<'p> {
             ast::Expr::Identifier(identifier) => {
                 self.hook.pre(self, super::ASTNode::Expr(key));
                 self.write_ast_comments(key);
-                // Judge contains string identifier, e.g., "x-y-z"
+                // Judge contains string or dot identifier, e.g., "x-y-z" and "a.b.c"
                 let names = &identifier.names;
 
                 let re = fancy_regex::Regex::new(IDENTIFIER_REGEX).unwrap();

kclvm/query/src/lib.rs

@@ -4,6 +4,7 @@
 //! input file name, and according to the ast::OverrideSpec transforms the nodes in the
 //! AST, recursively modifying or deleting the values of the nodes in the AST.
 pub mod r#override;
+pub mod path;
 pub mod query;
 pub mod selector;
 

kclvm/query/src/override.rs

@@ -12,6 +12,8 @@ use kclvm_ast_pretty::print_ast_module;
 use kclvm_parser::parse_expr;
 use kclvm_sema::pre_process::{fix_config_expr_nest_attr, transform_multi_assign};
 
+use crate::path::parse_attribute_path;
+
 use super::util::{invalid_spec_error, split_field_path};
 
 /// Import statement column offset always start with 1.
@@ -105,16 +107,15 @@ pub fn apply_override_on_module(
 ) -> Result<bool> {
     // Apply import paths on AST module.
     apply_import_paths_on_module(m, import_paths)?;
-    let ss = o.field_path.split('.').collect::<Vec<&str>>();
+    let ss = parse_attribute_path(&o.field_path)?;
     if ss.len() <= 1 {
         return Ok(false);
     }
-    let target_id = ss[0];
-    let field = ss[1..].join(".");
+    let target_id = &ss[0];
     let value = &o.field_value;
     let key = ast::Identifier {
-        names: field
-            .split('.')
+        names: ss[1..]
+            .iter()
             .map(|s| ast::Node::dummy_node(s.to_string()))
             .collect(),
         ctx: ast::ExprContext::Store,
@@ -132,7 +133,7 @@ pub fn apply_override_on_module(
     transform_multi_assign(m);
     let mut transformer = OverrideTransformer {
         target_id: target_id.to_string(),
-        field_path: field,
+        field_paths: ss[1..].to_vec(),
         override_key: key,
         override_value: build_expr_from_string(value),
         override_target_count: 0,
@@ -230,7 +231,7 @@ fn apply_import_paths_on_module(m: &mut ast::Module, import_paths: &[String]) ->
 /// OverrideTransformer is used to walk AST and transform it with the override values.
 struct OverrideTransformer {
     pub target_id: String,
-    pub field_path: String,
+    pub field_paths: Vec<String>,
     pub override_key: ast::Identifier,
     pub override_value: Option<ast::NodeRef<ast::Expr>>,
     pub override_target_count: usize,
@@ -324,7 +325,11 @@ impl OverrideTransformer {
     /// return whether is found a replaced one.
     fn lookup_config_and_replace(&self, config_expr: &mut ast::ConfigExpr) -> bool {
         // Split a path into multiple parts. `a.b.c` -> ["a", "b", "c"]
-        let parts = self.field_path.split('.').collect::<Vec<&str>>();
+        let parts = self
+            .field_paths
+            .iter()
+            .map(|s| s.as_str())
+            .collect::<Vec<&str>>();
         self.replace_config_with_path_parts(config_expr, &parts)
     }
 

kclvm/query/src/path.rs

@@ -0,0 +1,60 @@
+use anyhow::Result;
+
+/// Parse attribute path which returns either a vector of strings or an error. e.g.
+/// `a.b.c`, `a['b'].c`, `a["b"].c`, `a.['b'].c` and `a.["b"].c` both return `["a", "b", "c"]`
+pub fn parse_attribute_path(path: &str) -> Result<Vec<String>> {
+    let mut parts: Vec<String> = Vec::new();
+    let mut current = String::new();
+    let mut chars = path.chars().peekable();
+    let mut in_brackets = false;
+
+    while let Some(ch) = chars.next() {
+        if in_brackets {
+            if ch == '"' || ch == '\'' {
+                // Expecting the closing quote, skip if found
+                if chars.peek() == Some(&']') {
+                    chars.next(); // Consume the closing bracket
+                    in_brackets = false;
+                    continue;
+                }
+                return Err(anyhow::anyhow!("Expected closing bracket"));
+            } else {
+                current.push(ch);
+            }
+        } else {
+            match ch {
+                '.' => {
+                    if !current.is_empty() {
+                        parts.push(current.clone());
+                        current.clear();
+                    }
+                }
+                '[' => {
+                    if !current.is_empty() {
+                        parts.push(current.clone());
+                        current.clear();
+                    }
+                    in_brackets = true;
+                    // Skip the opening quote
+                    if let Some(next_char) = chars.next() {
+                        if next_char != '"' && next_char != '\'' {
+                            return Err(anyhow::anyhow!("Expected opening quote after '['"));
+                        }
+                    }
+                }
+                ']' => {
+                    return Err(anyhow::anyhow!("Unmatched closing bracket"));
+                }
+                _ => {
+                    current.push(ch);
+                }
+            }
+        }
+    }
+
+    if !current.is_empty() {
+        parts.push(current);
+    }
+
+    Ok(parts)
+}

kclvm/query/src/tests.rs

@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 
 use super::{r#override::apply_override_on_module, *};
+use crate::path::parse_attribute_path;
 use kclvm_ast::ast;
 use kclvm_parser::parse_file_force_errors;
 use pretty_assertions::assert_eq;
@@ -55,6 +56,7 @@ fn test_override_file_config() {
         "appConfiguration.mainContainer.name=override_name".to_string(),
         "appConfiguration.labels.key.key=\"override_value\"".to_string(),
         "appConfiguration.labels.key.str-key=\"override_value\"".to_string(),
+        "appConfiguration.labels.key['dot.key']=\"override_value\"".to_string(),
         "appConfiguration.overQuota=False".to_string(),
         "appConfiguration.probe={periodSeconds=20}".to_string(),
         "appConfiguration.resource-".to_string(),
@@ -109,6 +111,7 @@ appConfiguration = AppConfiguration {
         key: {
             key: "override_value"
             "str-key" = "override_value"
+            "dot.key" = "override_value"
         }
     }
     mainContainer: Main {name: "override_name"}
@@ -142,3 +145,42 @@ fn test_parse_override_spec_invalid() {
         assert!(parse_override_spec(spec).is_err(), "{spec} test failed");
     }
 }
+
+#[test]
+fn test_parse_property_path() {
+    assert_eq!(parse_attribute_path("a.b.c").unwrap(), vec!["a", "b", "c"]);
+    assert_eq!(
+        parse_attribute_path(r#"a["b"].c"#).unwrap(),
+        vec!["a", "b", "c"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.["b"].c"#).unwrap(),
+        vec!["a", "b", "c"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a['b'].c"#).unwrap(),
+        vec!["a", "b", "c"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.b['c.d']"#).unwrap(),
+        vec!["a", "b", "c.d"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.b.['c.d']"#).unwrap(),
+        vec!["a", "b", "c.d"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.b['c.d'].e"#).unwrap(),
+        vec!["a", "b", "c.d", "e"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.b.['c.d'].e"#).unwrap(),
+        vec!["a", "b", "c.d", "e"]
+    );
+    assert_eq!(
+        parse_attribute_path(r#"a.b.c-d.e"#).unwrap(),
+        vec!["a", "b", "c-d", "e"]
+    );
+    assert!(parse_attribute_path(r#"a.[b.c-d.e"#).is_err(),);
+    assert!(parse_attribute_path(r#"a.[b.c]-d.e"#).is_err(),);
+}

kclvm/tools/src/LSP/src/rename.rs

@@ -8,7 +8,7 @@ use chumsky::chain::Chain;
 use kclvm_ast::ast::{self, Program};
 use kclvm_error::diagnostic;
 use kclvm_parser::{load_program, LoadProgramOptions, ParseSessionRef};
-use kclvm_query::selector::parse_symbol_selector_spec;
+use kclvm_query::{path::parse_attribute_path, selector::parse_symbol_selector_spec};
 use kclvm_sema::{
     advanced_resolver::AdvancedResolver, core::global_state::GlobalState, namer::Namer,
     resolver::resolve_program_with_opts,
@@ -120,7 +120,7 @@ where
     F: Fn(String) -> VfsPath,
 {
     let mut pkg = PathBuf::from(&symbol_spec.pkg_root);
-    let fields: Vec<&str> = symbol_spec.field_path.split(".").collect();
+    let fields = parse_attribute_path(&symbol_spec.field_path).unwrap_or_default();
     if !symbol_spec.pkgpath.is_empty() {
         let pkg_names = symbol_spec.pkgpath.split(".");
         for n in pkg_names {
@@ -143,7 +143,7 @@ where
         {
             let mut owner_ref = symbol_ref;
             let mut target = None;
-            for field in fields {
+            for field in &fields {
                 let owner = gs.get_symbols().get_symbol(owner_ref).unwrap();
                 target = owner.get_attribute(field, gs.get_symbols(), None);
                 if let Some(target) = target {