Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: override spec with more attribute config operation including : and += #1358

Merged
merged 1 commit into from
May 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions kclvm/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion kclvm/ast/src/ast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -341,10 +341,10 @@ pub struct CmdArgSpec {
/// KCL command line override spec, e.g. `kcl main.k -O pkgpath:path.to.field=field_value`
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct OverrideSpec {
pub pkgpath: String,
pub field_path: String,
pub field_value: String,
pub action: OverrideAction,
pub operation: ConfigEntryOperation,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
Expand Down
1 change: 1 addition & 0 deletions kclvm/query/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ kclvm-sema = {path = "../sema"}
kclvm-error = {path = "../error"}
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
fancy-regex = "0.7.1"
maplit = "1.0.2"

[dev-dependencies]
Expand Down
12 changes: 2 additions & 10 deletions kclvm/query/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ mod tests;
mod util;

use anyhow::{anyhow, Result};
use kclvm_ast::ast;
use kclvm_ast_pretty::print_ast_module;
use kclvm_error::diagnostic::Errors;
use kclvm_parser::parse_file;
Expand All @@ -23,8 +22,6 @@ use kclvm_sema::pre_process::fix_config_expr_nest_attr;
pub use query::{get_schema_type, GetSchemaOption};
pub use r#override::{apply_override_on_module, apply_overrides};

use self::r#override::parse_override_spec;

/// Override and rewrite a file with override specifications. Please note that this is an external user API,
/// and it can directly modify the KCL file in place.
///
Expand Down Expand Up @@ -84,20 +81,15 @@ pub fn override_file(
specs: &[String],
import_paths: &[String],
) -> Result<OverrideFileResult> {
// Parse override spec strings.
let overrides = specs
.iter()
.map(|s| parse_override_spec(s))
.collect::<Result<Vec<ast::OverrideSpec>, _>>()?;
// Parse file to AST module.
let mut parse_result = match parse_file(file, None) {
Ok(module) => module,
Err(msg) => return Err(anyhow!("{}", msg)),
};
let mut result = false;
// Override AST module.
for o in &overrides {
if apply_override_on_module(&mut parse_result.module, o, import_paths)? {
for s in specs {
if apply_override_on_module(&mut parse_result.module, s, import_paths)? {
result = true;
}
}
Expand Down
210 changes: 163 additions & 47 deletions kclvm/query/src/override.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use kclvm_sema::pre_process::{fix_config_expr_nest_attr, transform_multi_assign}

use crate::{node::AstNodeMover, path::parse_attribute_path};

use super::util::{invalid_spec_error, split_field_path};
use super::util::invalid_spec_error;

/// Import statement column offset always start with 1.
/// todo: The (1-based) column offset needs to be constrained by specifications.
Expand All @@ -40,17 +40,12 @@ const IMPORT_STMT_COLUMN_OFFSET: u64 = 1;
/// ```
pub fn apply_overrides(
prog: &mut ast::Program,
overrides: &[ast::OverrideSpec],
overrides: &[String],
import_paths: &[String],
print_ast: bool,
) -> Result<()> {
for o in overrides {
let pkgpath = if o.pkgpath.is_empty() {
MAIN_PKG
} else {
&o.pkgpath
};
if let Some(modules) = prog.pkgs.get_mut(pkgpath) {
if let Some(modules) = prog.pkgs.get_mut(MAIN_PKG) {
for m in modules.iter_mut() {
if apply_override_on_module(m, o, import_paths)? && print_ast {
let code_str = print_ast_module(m);
Expand Down Expand Up @@ -103,11 +98,12 @@ fn build_expr_from_string(value: &str) -> Option<ast::NodeRef<ast::Expr>> {
/// ```
pub fn apply_override_on_module(
m: &mut ast::Module,
o: &ast::OverrideSpec,
o: &str,
import_paths: &[String],
) -> Result<bool> {
// Apply import paths on AST module.
apply_import_paths_on_module(m, import_paths)?;
let o = parse_override_spec(o)?;
let ss = parse_attribute_path(&o.field_path)?;
let default = String::default();
let target_id = ss.get(0).unwrap_or(&default);
Expand Down Expand Up @@ -137,7 +133,8 @@ pub fn apply_override_on_module(
override_value: build_expr_from_string(value),
override_target_count: 0,
has_override: false,
action: o.action.clone(),
action: o.action,
operation: o.operation,
};
transformer.walk_module(m);
Ok(transformer.has_override)
Expand All @@ -152,36 +149,97 @@ pub fn apply_override_on_module(
/// action: ast::OverrideAction::CreateOrUpdate,
/// }
pub fn parse_override_spec(spec: &str) -> Result<ast::OverrideSpec> {
if spec.contains('=') {
if let Some((path, value, operation)) = split_override_spec_op(spec) {
// Create or update the override value.
let split_values = spec.splitn(2, '=').map(|s| s.trim()).collect::<Vec<&str>>();
let path = split_values
.first()
.ok_or_else(|| invalid_spec_error(spec))?;
let field_value = split_values
.get(1)
.ok_or_else(|| invalid_spec_error(spec))?;
let (pkgpath, field_path) = split_field_path(path)?;
Ok(ast::OverrideSpec {
pkgpath,
field_path,
field_value: field_value.to_string(),
action: ast::OverrideAction::CreateOrUpdate,
})
let field_path = path.trim().to_string();
let field_value = value.trim().to_string();
if field_path.is_empty() || field_value.is_empty() {
Err(invalid_spec_error(spec))
} else {
Ok(ast::OverrideSpec {
field_path,
field_value,
action: ast::OverrideAction::CreateOrUpdate,
operation,
})
}
} else if let Some(stripped_spec) = spec.strip_suffix('-') {
// Delete the override value.
let (pkgpath, field_path) = split_field_path(stripped_spec)?;
Ok(ast::OverrideSpec {
pkgpath,
field_path,
field_value: "".to_string(),
action: ast::OverrideAction::Delete,
})
let field_path = stripped_spec.trim().to_string();
if field_path.is_empty() {
Err(invalid_spec_error(spec))
} else {
Ok(ast::OverrideSpec {
field_path: stripped_spec.trim().to_string(),
field_value: "".to_string(),
action: ast::OverrideAction::Delete,
operation: ast::ConfigEntryOperation::Override,
})
}
} else {
Err(invalid_spec_error(spec))
}
}

/// split_override_spec_op split the override_spec and do not split the override_op in list
/// expr, dict expr and string e.g., "a.b=1" -> (a.b, 1, =), "a["a=1"]=1" -> (a["a=1"], =, 1)
pub fn split_override_spec_op(spec: &str) -> Option<(String, String, ast::ConfigEntryOperation)> {
let mut i = 0;
let mut stack = String::new();
while i < spec.chars().count() {
let (c_idx, c) = spec.char_indices().nth(i).unwrap();
if c == '=' && stack.is_empty() {
return Some((
spec[..c_idx].to_string(),
spec[c_idx + 1..].to_string(),
ast::ConfigEntryOperation::Override,
));
} else if c == ':' && stack.is_empty() {
return Some((
spec[..c_idx].to_string(),
spec[c_idx + 1..].to_string(),
ast::ConfigEntryOperation::Union,
));
} else if c == '+' && stack.is_empty() {
if let Some((c_next_idx, c_next)) = spec.char_indices().nth(i + 1) {
if c_next == '=' {
return Some((
spec[..c_idx].to_string(),
spec[c_next_idx + 1..].to_string(),
ast::ConfigEntryOperation::Insert,
));
}
}
}
// List/Dict type
else if c == '[' || c == '{' {
stack.push(c);
}
// List/Dict type
else if c == ']' || c == '}' {
stack.pop();
}
// String literal type
else if c == '\"' {
let t: String = spec.chars().skip(i).collect();
let re = fancy_regex::Regex::new(r#""(?!"").*?(?<!\\)(\\\\)*?""#).unwrap();
if let Ok(Some(v)) = re.find(&t) {
i += v.as_str().chars().count() - 1;
}
}
// String literal type
else if c == '\'' {
let t: String = spec.chars().skip(i).collect();
let re = fancy_regex::Regex::new(r#"'(?!'').*?(?<!\\)(\\\\)*?'"#).unwrap();
if let Ok(Some(v)) = re.find(&t) {
i += v.as_str().chars().count() - 1;
}
}
i += 1;
}
None
}

// Transform the AST module with the import path list.
fn apply_import_paths_on_module(m: &mut ast::Module, import_paths: &[String]) -> Result<()> {
if import_paths.is_empty() {
Expand Down Expand Up @@ -248,6 +306,7 @@ struct OverrideTransformer {
pub override_target_count: usize,
pub has_override: bool,
pub action: ast::OverrideAction,
pub operation: ast::ConfigEntryOperation,
}

impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer {
Expand Down Expand Up @@ -365,24 +424,81 @@ impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer {
},
)))),
value: self.clone_override_value(),
operation: ast::ConfigEntryOperation::Override,
operation: self.operation.clone(),
insert_index: -1,
}))],
})))
};
match &self.operation {
ast::ConfigEntryOperation::Override => {
let assign = ast::AssignStmt {
targets: vec![Box::new(ast::Node::dummy_node(ast::Identifier {
names: vec![ast::Node::dummy_node(self.target_id.clone())],
ctx: ast::ExprContext::Store,
pkgpath: "".to_string(),
}))],
ty: None,
value,
};
module
.body
.push(Box::new(ast::Node::dummy_node(ast::Stmt::Assign(assign))));
}
ast::ConfigEntryOperation::Union => {
let schema_expr: Result<ast::Node<ast::SchemaExpr>, _> =
value.as_ref().clone().try_into();
match schema_expr {
Ok(schema_expr) => {
let stmt = ast::UnificationStmt {
target: Box::new(ast::Node::dummy_node(ast::Identifier {
names: vec![ast::Node::dummy_node(
self.target_id.clone(),
)],
ctx: ast::ExprContext::Store,
pkgpath: "".to_string(),
})),
value: Box::new(schema_expr),
};
module.body.push(Box::new(ast::Node::dummy_node(
ast::Stmt::Unification(stmt),
)));
}
Err(_) => {
let stmt = ast::AssignStmt {
targets: vec![Box::new(ast::Node::dummy_node(
ast::Identifier {
names: vec![ast::Node::dummy_node(
self.target_id.clone(),
)],
ctx: ast::ExprContext::Store,
pkgpath: "".to_string(),
},
))],
ty: None,
value,
};
module.body.push(Box::new(ast::Node::dummy_node(
ast::Stmt::Assign(stmt),
)));
}
}
}
ast::ConfigEntryOperation::Insert => {
let stmt = ast::AugAssignStmt {
target: Box::new(ast::Node::dummy_node(ast::Identifier {
names: vec![ast::Node::dummy_node(self.target_id.clone())],
ctx: ast::ExprContext::Store,
pkgpath: "".to_string(),
})),
op: ast::AugOp::Add,
value,
};
module
.body
.push(Box::new(ast::Node::dummy_node(ast::Stmt::AugAssign(stmt))));
}
}

let assign = ast::AssignStmt {
targets: vec![Box::new(ast::Node::dummy_node(ast::Identifier {
names: vec![ast::Node::dummy_node(self.target_id.clone())],
ctx: ast::ExprContext::Store,
pkgpath: "".to_string(),
}))],
ty: None,
value,
};
module
.body
.push(Box::new(ast::Node::dummy_node(ast::Stmt::Assign(assign))));
self.has_override = true;
}
ast::OverrideAction::Delete => {
Expand Down Expand Up @@ -451,7 +567,7 @@ impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer {
self.override_key.clone(),
)))),
value: self.clone_override_value(),
operation: ast::ConfigEntryOperation::Override,
operation: self.operation.clone(),
insert_index: -1,
})));
self.has_override = true;
Expand Down Expand Up @@ -585,7 +701,7 @@ impl OverrideTransformer {
.push(Box::new(ast::Node::dummy_node(ast::ConfigEntry {
key: Some(Box::new(ast::Node::dummy_node(ast::Expr::Identifier(key)))),
value: self.clone_override_value(),
operation: ast::ConfigEntryOperation::Override,
operation: self.operation.clone(),
insert_index: -1,
})));
changed = true;
Expand Down
Loading
Loading