Skip to content

Commit

Permalink
feat: add type name for schema_expr in list_variables (#1279)
Browse files Browse the repository at this point in the history 
Signed-off-by: zongz <zongzhe1024@163.com>
  • Loading branch information
@zong-zhe
zong-zhe committed May 6, 2024
1 parent e028d47 commit 4a62bd2
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 45 deletions.
5 changes: 3 additions & 2 deletions kclvm/api/src/service/service_impl.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -349,11 +349,12 @@ impl KclvmServiceImpl {
let variables: HashMap<String, Variable> = select_res
.select_result
.iter()
.map(|(key, value)| {
.map(|(key, var)| {
(
key.clone(),
Variable {
value: value.to_string(),
value: var.value.to_string(),
type_name: var.type_name.to_string(),
},
)
})
Expand Down
55 changes: 47 additions & 8 deletions kclvm/query/src/selector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ pub struct UnsupportedSelectee {
/// Selector is used to select the target variable from the kcl program.
pub struct Selector {
select_specs: Vec<String>,
select_result: HashMap<String, String>,
select_result: HashMap<String, Variable>,
unsupported: Vec<UnsupportedSelectee>,
inner: SelectorInner,
}
Expand Down Expand Up @@ -147,7 +147,11 @@ impl<'ctx> MutSelfWalker for Selector {
let kcode = print_ast_node(ASTNode::Expr(&Box::new(ast::Node::dummy_node(
ast::Expr::Schema(unification_stmt.value.node.clone()),
))));
self.select_result.insert(target.to_string(), kcode);

self.select_result.insert(
target.to_string(),
Variable::new(unification_stmt.value.node.name.node.get_name(), kcode),
);
} else {
// if length of spec is largr or equal to target
let selector = self.inner.pop_front();
Expand All @@ -159,7 +163,10 @@ impl<'ctx> MutSelfWalker for Selector {
print_ast_node(ASTNode::Expr(&Box::new(ast::Node::dummy_node(
ast::Expr::Schema(unification_stmt.value.node.clone()),
))));
self.select_result.insert(target.to_string(), kcode);
self.select_result.insert(
target.to_string(),
Variable::new(unification_stmt.value.node.name.node.get_name(), kcode),
);
} else {
// walk ahead
self.walk_schema_expr(&unification_stmt.value.node);
Expand All @@ -182,14 +189,20 @@ impl<'ctx> MutSelfWalker for Selector {
}
// get the value source code of the assign statement
let kcode = print_ast_node(ASTNode::Expr(&assign_stmt.value));
let type_name = if let ast::Expr::Schema(schema) = &assign_stmt.value.node {
schema.name.node.get_name()
} else {
"".to_string()
};
// The length of name for variable in top level is 1
if assign_stmt.targets.len() == 1 {
let target = &assign_stmt.targets[0];
let target = &Some(Box::new(ast::Node::dummy_node(ast::Expr::Identifier(
target.node.clone(),
))));
let key = get_key_path(&target);
self.select_result.insert(key.to_string(), kcode);
self.select_result
.insert(key.to_string(), Variable::new(type_name, kcode));
}
} else {
// Compare the target with the spec
Expand All @@ -211,7 +224,14 @@ impl<'ctx> MutSelfWalker for Selector {

// matched
let kcode = print_ast_node(ASTNode::Expr(&assign_stmt.value));
self.select_result.insert(target.to_string(), kcode);
let type_name =
if let ast::Expr::Schema(schema) = &assign_stmt.value.node {
schema.name.node.get_name()
} else {
"".to_string()
};
self.select_result
.insert(target.to_string(), Variable::new(type_name, kcode));
} else {
// walk ahead
self.walk_expr(&assign_stmt.value.node)
Expand Down Expand Up @@ -249,8 +269,15 @@ impl<'ctx> MutSelfWalker for Selector {
return;
}
let kcode = print_ast_node(ASTNode::Expr(&item.node.value));
self.select_result
.insert(self.inner.current_spec.to_string(), kcode);
let type_name = if let ast::Expr::Schema(schema) = &item.node.value.node {
schema.name.node.get_name()
} else {
"".to_string()
};
self.select_result.insert(
self.inner.current_spec.to_string(),
Variable::new(type_name, kcode),
);
} else {
// the spec is still not used up
// walk ahead
Expand Down Expand Up @@ -333,10 +360,22 @@ impl<'ctx> MutSelfWalker for Selector {
}

pub struct ListVariablesResult {
pub select_result: HashMap<String, String>,
pub select_result: HashMap<String, Variable>,
pub unsupported: Vec<UnsupportedSelectee>,
}

#[derive(Debug, PartialEq)]
pub struct Variable {
pub type_name: String,
pub value: String,
}

impl Variable {
pub fn new(type_name: String, value: String) -> Self {
Self { type_name, value }
}
}

/// list_options provides users with the ability to parse kcl program and get all option
/// calling information.
pub fn list_variables(file: String, specs: Vec<String>) -> Result<ListVariablesResult> {
Expand Down
4 changes: 4 additions & 0 deletions kclvm/query/src/test_data/test_list_variables/supported.k
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,4 +76,8 @@ schema Job:

job = Job {
name = "{}-{}".format("app", "test").lower()
}

select = a.b.c {
a: 1
}
92 changes: 57 additions & 35 deletions kclvm/query/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,16 +225,16 @@ fn test_list_variables() {
.display()
.to_string();
let test_cases = vec![
("a", "1"),
("a1", "2"),
("a3", "3m"),
("b1", "True"),
("b2", "False"),
("s1", "\"Hello\""),
("array1", "[1, 2, 3]"),
("dict1", "{\"a\": 1, \"b\": 2}"),
("dict1.a", "1"),
("dict1.b", "2"),
("a", "1", ""),
("a1", "2", ""),
("a3", "3m", ""),
("b1", "True", ""),
("b2", "False", ""),
("s1", "\"Hello\"", ""),
("array1", "[1, 2, 3]", ""),
("dict1", "{\"a\": 1, \"b\": 2}", ""),
("dict1.a", "1", ""),
("dict1.b", "2", ""),
(
"dict2",
r#"{
Expand All @@ -244,17 +244,19 @@ fn test_list_variables() {
"d": 3
}
}"#,
"",
),
("dict2.a", "1"),
("dict2.a", "1", ""),
(
"dict2.b",
r#"{
"c": 2
"d": 3
}"#,
"",
),
("dict2.b.c", "2"),
("dict2.b.d", "3"),
("dict2.b.c", "2", ""),
("dict2.b.d", "3", ""),
(
"sha",
r#"A {
Expand All @@ -266,25 +268,28 @@ fn test_list_variables() {
}
}
}"#,
"A",
),
("sha.name", "\"Hello\""),
("sha.ids", "[1, 2, 3]"),
("sha.name", "\"Hello\"", ""),
("sha.ids", "[1, 2, 3]", ""),
(
"sha.data",
r#"{
"a": {
"b": {"c": 2}
}
}"#,
"",
),
(
"sha.data.a",
r#"{
"b": {"c": 2}
}"#,
"",
),
("sha.data.a.b", r#"{"c": 2}"#),
("sha.data.a.b.c", "2"),
("sha.data.a.b", r#"{"c": 2}"#, ""),
("sha.data.a.b.c", "2", ""),
(
"shb",
r#"B {
Expand All @@ -298,6 +303,7 @@ fn test_list_variables() {
}
}
}"#,
"B",
),
(
"shb.a",
Expand All @@ -310,33 +316,40 @@ fn test_list_variables() {
}
}
}"#,
"",
),
("shb.a.name", "\"HelloB\""),
("shb.a.ids", "[4, 5, 6]"),
("shb.a.name", "\"HelloB\"", ""),
("shb.a.ids", "[4, 5, 6]", ""),
(
"shb.a.data",
r#"{
"d": {
"e": {"f": 3}
}
}"#,
"",
),
(
"shb.a.data.d",
r#"{
"e": {"f": 3}
}"#,
"",
),
("shb.a.data.d.e", "{\"f\": 3}"),
("uconfa.name", "\"b\""),
("c.a", "{ids: [7, 8, 9]}"),
("job.name", r#""{}-{}".format("app", "test").lower()"#),
("shb.a.data.d.e", "{\"f\": 3}", ""),
("uconfa.name", "\"b\"", ""),
("c.a", "{ids: [7, 8, 9]}", ""),
("job.name", r#""{}-{}".format("app", "test").lower()"#, ""),
];

for (spec, expected) in test_cases {
for (spec, expected, expected_name) in test_cases {
let specs = vec![spec.to_string()];
let result = list_variables(file.clone(), specs).unwrap();
assert_eq!(result.select_result.get(spec).unwrap(), expected);
assert_eq!(result.select_result.get(spec).unwrap().value, expected);
assert_eq!(
result.select_result.get(spec).unwrap().type_name,
expected_name
);
}
}

Expand All @@ -348,14 +361,14 @@ fn test_list_all_variables() {
.display()
.to_string();
let test_cases = vec![
("a", "1"),
("a1", "2"),
("a3", "3m"),
("b1", "True"),
("b2", "False"),
("s1", "\"Hello\""),
("array1", "[1, 2, 3]"),
("dict1", "{\"a\": 1, \"b\": 2}"),
("a", "1", ""),
("a1", "2", ""),
("a3", "3m", ""),
("b1", "True", ""),
("b2", "False", ""),
("s1", "\"Hello\"", ""),
("array1", "[1, 2, 3]", ""),
("dict1", "{\"a\": 1, \"b\": 2}", ""),
(
"dict2",
r#"{
Expand All @@ -365,6 +378,7 @@ fn test_list_all_variables() {
"d": 3
}
}"#,
"",
),
(
"sha",
Expand All @@ -377,6 +391,7 @@ fn test_list_all_variables() {
}
}
}"#,
"A",
),
(
"shb",
Expand All @@ -391,16 +406,23 @@ fn test_list_all_variables() {
}
}
}"#,
"B",
),
(
"job",
r#"Job {name = "{}-{}".format("app", "test").lower()}"#,
"Job",
),
("select", r#"a.b.c {a: 1}"#, "a.b.c"),
];

for (spec, expected) in test_cases {
for (spec, expected, expected_name) in test_cases {
let result = list_variables(file.clone(), vec![]).unwrap();
assert_eq!(result.select_result.get(spec).unwrap(), expected);
assert_eq!(result.select_result.get(spec).unwrap().value, expected);
assert_eq!(
result.select_result.get(spec).unwrap().type_name,
expected_name
);
}
}

Expand Down
1 change: 1 addition & 0 deletions kclvm/spec/gpyrpc/gpyrpc.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,7 @@ message ListVariables_Result {

message Variable {
string value = 1;
string type_name = 2;
}

message GetFullSchemaType_Args {
Expand Down

0 comments on commit 4a62bd2

Please sign in to comment.