feat: add service api GetSchemaTypeMappingUnderPath (#1766)

* feat: add service api GetSchemaTypeMappingUnderPath

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

* fix typo

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

---------

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

kclvm/api/src/service/capi.rs

@@ -171,6 +171,9 @@ pub(crate) fn kclvm_get_service_fn_ptr_by_name(name: &str) -> u64 {
         "KclvmService.ExecArtifact" => exec_artifact as *const () as u64,
         "KclvmService.OverrideFile" => override_file as *const () as u64,
         "KclvmService.GetSchemaTypeMapping" => get_schema_type_mapping as *const () as u64,
+        "KclvmService.GetSchemaTypeMappingUnderPath" => {
+            get_schema_type_mapping_under_path as *const () as u64
+        }
         "KclvmService.FormatCode" => format_code as *const () as u64,
         "KclvmService.FormatPath" => format_path as *const () as u64,
         "KclvmService.LintPath" => lint_path as *const () as u64,
@@ -523,6 +526,30 @@ pub(crate) fn get_schema_type_mapping(
     )
 }
 
+/// Get schema types under path
+///
+/// # Parameters
+/// file: [&str]. The kcl filename.
+///
+/// code: [Option<&str>]. The kcl code string
+///
+/// schema_name: [Option<&str>]. The schema name, when the schema name is empty, all schemas are returned.
+pub(crate) fn get_schema_type_mapping_under_path(
+    serv: *mut kclvm_service,
+    args: *const c_char,
+    args_len: usize,
+    result_len: *mut usize,
+) -> *const c_char {
+    call!(
+        serv,
+        args,
+        args_len,
+        result_len,
+        GetSchemaTypeMappingArgs,
+        get_schema_type_mapping_under_path
+    )
+}
+
 /// Service for formatting a code source and returns the formatted source and
 /// whether the source is changed.
 pub(crate) fn format_code(

kclvm/api/src/service/jsonrpc.rs

@@ -157,6 +157,21 @@ fn register_kclvm_service(io: &mut IoHandler) {
         };
         futures::future::ready(catch!(kclvm_service_impl, args, get_schema_type_mapping))
     });
+    io.add_method(
+        "KclvmService.GetSchemaTypeMappingUnderPath",
+        |params: Params| {
+            let kclvm_service_impl = KclvmServiceImpl::default();
+            let args: GetSchemaTypeMappingArgs = match params.parse() {
+                Ok(val) => val,
+                Err(err) => return futures::future::ready(Err(err)),
+            };
+            futures::future::ready(catch!(
+                kclvm_service_impl,
+                args,
+                get_schema_type_mapping_under_path
+            ))
+        },
+    );
     io.add_method("KclvmService.FormatCode", |params: Params| {
         let kclvm_service_impl = KclvmServiceImpl::default();
         let args: FormatCodeArgs = match params.parse() {

kclvm/api/src/service/service_impl.rs

@@ -3,7 +3,7 @@ use std::io::Write;
 use std::path::PathBuf;
 use std::string::String;
 
-use crate::gpyrpc::*;
+use crate::gpyrpc::{self, *};
 
 use kcl_language_server::rename;
 use kclvm_ast::ast::SerializeProgram;
@@ -17,8 +17,8 @@ use kclvm_parser::KCLModuleCache;
 use kclvm_parser::LoadProgramOptions;
 use kclvm_parser::ParseSessionRef;
 use kclvm_query::override_file;
-use kclvm_query::query::get_full_schema_type;
 use kclvm_query::query::CompilationOptions;
+use kclvm_query::query::{get_full_schema_type, get_full_schema_type_under_path};
 use kclvm_query::selector::{list_variables, ListOptions};
 use kclvm_query::GetSchemaOption;
 use kclvm_runner::exec_program;
@@ -666,6 +666,74 @@ impl KclvmServiceImpl {
         })
     }
 
+    /// Service for getting the schema mapping under path.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kclvm_api::service::service_impl::KclvmServiceImpl;
+    /// use kclvm_api::gpyrpc::*;
+    /// use std::path::Path;
+    /// use kclvm_ast::MAIN_PKG;
+    ///
+    /// let serv = KclvmServiceImpl::default();
+    /// let work_dir_parent = Path::new(".").join("src").join("testdata").join("get_schema_ty_under_path");
+    /// let args = ExecProgramArgs {
+    ///     k_filename_list: vec![
+    ///         work_dir_parent.join("aaa").canonicalize().unwrap().display().to_string()
+    ///     ],
+    ///     external_pkgs: vec![
+    ///         ExternalPkg {
+    ///             pkg_name:"bbb".to_string(),
+    ///             pkg_path: work_dir_parent.join("bbb").canonicalize().unwrap().display().to_string()
+    ///         },
+    ///         ExternalPkg {
+    ///             pkg_name:"helloworld".to_string(),
+    ///             pkg_path: work_dir_parent.join("helloworld_0.0.1").canonicalize().unwrap().display().to_string()
+    ///         },
+    ///     ],
+    ///     ..Default::default()
+    /// };
+    ///
+    /// let result = serv.get_schema_type_mapping_under_path(&GetSchemaTypeMappingArgs {
+    ///     exec_args: Some(args),
+    ///     ..Default::default()
+    /// }).unwrap();
+    ///  assert_eq!(result.schema_type_mapping.get(MAIN_PKG).unwrap().schema_type.len(), 1);
+    ///  assert_eq!(result.schema_type_mapping.get("bbb").unwrap().schema_type.len(), 2);
+    ///  assert_eq!(result.schema_type_mapping.get("helloworld").unwrap().schema_type.len(), 1);
+    ///  assert_eq!(result.schema_type_mapping.get("sub").unwrap().schema_type.len(), 1);
+    /// ```
+    pub fn get_schema_type_mapping_under_path(
+        &self,
+        args: &GetSchemaTypeMappingArgs,
+    ) -> anyhow::Result<GetSchemaTypeMappingUnderPathResult> {
+        let mut type_mapping = HashMap::new();
+        let exec_args = transform_exec_para(&args.exec_args, self.plugin_agent)?;
+        for (k, schema_tys) in get_full_schema_type_under_path(
+            Some(&args.schema_name),
+            CompilationOptions {
+                paths: exec_args.clone().k_filename_list,
+                loader_opts: Some(exec_args.get_load_program_options()),
+                resolve_opts: Options {
+                    resolve_val: true,
+                    ..Default::default()
+                },
+                get_schema_opts: GetSchemaOption::Definitions,
+            },
+        )? {
+            let mut tys = vec![];
+            for schema_ty in schema_tys {
+                tys.push(kcl_schema_ty_to_pb_ty(&schema_ty));
+            }
+            type_mapping.insert(k, gpyrpc::SchemaTypes { schema_type: tys });
+        }
+
+        Ok(GetSchemaTypeMappingUnderPathResult {
+            schema_type_mapping: type_mapping,
+        })
+    }
+
     /// Service for formatting a code source and returns the formatted source and
     /// whether the source is changed.
     ///

kclvm/api/src/testdata/get_schema_ty_under_path/aaa/kcl.mod

@@ -0,0 +1,8 @@
+[package]
+name = "aaa"
+edition = "0.0.1"
+version = "0.0.1"
+
+[dependencies]
+bbb = { path = "../bbb" }
+helloworld = "0.0.1"

kclvm/api/src/testdata/get_schema_ty_under_path/aaa/main.k

@@ -0,0 +1,8 @@
+import bbb
+
+schema A:
+    name: str
+
+a = bbb.B {
+    name: "b instance in a"
+}

kclvm/api/src/testdata/get_schema_ty_under_path/aaa/sub/sub.k

@@ -0,0 +1,3 @@
+schema Sub:
+    name: str
+

kclvm/api/src/testdata/get_schema_ty_under_path/bbb/kcl.mod

@@ -0,0 +1,5 @@
+[package]
+name = "bbb"
+edition = "0.0.1"
+version = "0.0.1"
+

kclvm/api/src/testdata/get_schema_ty_under_path/bbb/main.k

@@ -0,0 +1,4 @@
+schema Base:
+    n: str
+schema B(Base):
+    name: str

kclvm/api/src/testdata/get_schema_ty_under_path/helloworld_0.0.1/README.md

@@ -0,0 +1,2 @@
+## Introduction
+This is a kcl package named helloworld.

kclvm/api/src/testdata/get_schema_ty_under_path/helloworld_0.0.1/kcl.mod

@@ -0,0 +1,5 @@
+[package]
+name = "helloworld"
+edition = "0.0.1"
+version = "0.0.1"
+

kclvm/api/src/testdata/get_schema_ty_under_path/helloworld_0.0.1/main.k

@@ -0,0 +1,3 @@
+The_first_kcl_program = 'Hello World!'
+schema Hello:
+    name: str

kclvm/spec/gpyrpc/gpyrpc.proto

@@ -1184,6 +1184,17 @@ message GetSchemaTypeMapping_Result {
 	map<string, KclType> schema_type_mapping = 1;
 }
 
+// Message for get schema type mapping response.
+message GetSchemaTypeMappingUnderPath_Result {
+	// Map of pkg and schema types mappings.
+	map<string, SchemaTypes> schema_type_mapping = 1;
+}
+
+message SchemaTypes {
+	// List of schema type mappings.
+	repeated KclType schema_type = 1;
+}
+
 // Message for validate code request arguments.
 message ValidateCode_Args {
 	// Path to the data file.