add return value for union type to get_type_symbol

Signed-off-by: shruti2522 <[email protected]>
kcl-lang · Jun 12, 2024 · d86a0a9 · d86a0a9
1 parent 19efb90
commit d86a0a9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 44 deletions.
diff --git a/kclvm/sema/src/core/symbol.rs b/kclvm/sema/src/core/symbol.rs
@@ -329,8 +329,16 @@ impl SymbolData {
             TypeKind::Dict(_) => None,
             TypeKind::NumberMultiplier(_) => None,
             TypeKind::Function(_) => None,
-            TypeKind::Union(_) => None,
-
+            TypeKind::Union(types) => {
+                if types
+                    .iter()
+                    .all(|ut| matches!(&ut.kind, TypeKind::StrLit(_) | TypeKind::Str))
+                {
+                    self.get_symbol_by_fully_qualified_name(BUILTIN_STR_PACKAGE)
+                } else {
+                    None
+                }
+            }
             TypeKind::Schema(schema_ty) => {
                 let fully_qualified_ty_name = schema_ty.pkgpath.clone() + "." + &schema_ty.name;
 
@@ -373,16 +381,7 @@ impl SymbolData {
             TypeKind::IntLit(_) => vec![],
             TypeKind::Float => vec![],
             TypeKind::FloatLit(_) => vec![],
-            TypeKind::Str => {
-                let mut result = vec![];
-                if let Some(symbol_ref) = self.get_type_symbol(ty, module_info) {
-                    if let Some(symbol) = self.get_symbol(symbol_ref) {
-                        result = symbol.get_all_attributes(self, module_info);
-                    }
-                }
-                result
-            }
-            TypeKind::StrLit(_) => {
+            TypeKind::Str | TypeKind::StrLit(_) => {
                 let mut result = vec![];
                 if let Some(symbol_ref) = self.get_type_symbol(ty, module_info) {
                     if let Some(symbol) = self.get_symbol(symbol_ref) {

diff --git a/kclvm/tools/src/LSP/src/completion.rs b/kclvm/tools/src/LSP/src/completion.rs
@@ -255,25 +255,7 @@ fn completion_dot(
         Some(def_ref) => {
             if let Some(def) = gs.get_symbols().get_symbol(def_ref) {
                 let module_info = gs.get_packages().get_module_info(&pos.filename);
-                let attrs = match &def.get_sema_info().ty {
-                    Some(symbol_ty) => match &symbol_ty.kind {
-                        TypeKind::Union(union_types) => {
-                            let mut union_attrs = vec![];
-                            for union_ty in union_types {
-                                if let TypeKind::StrLit(_) | TypeKind::Str = &union_ty.kind {
-                                    union_attrs.extend(gs.get_symbols().get_type_all_attribute(
-                                        union_ty,
-                                        "",
-                                        module_info,
-                                    ));
-                                }
-                            }
-                            union_attrs
-                        }
-                        _ => def.get_all_attributes(gs.get_symbols(), module_info),
-                    },
-                    None => def.get_all_attributes(gs.get_symbols(), module_info),
-                };
+                let attrs = def.get_all_attributes(gs.get_symbols(), module_info);
                 for attr in attrs {
                     let attr_def = gs.get_symbols().get_symbol(attr);
                     if let Some(attr_def) = attr_def {
@@ -295,18 +277,6 @@ fn completion_dot(
                                 };
                                 let kind = match &def.get_sema_info().ty {
                                     Some(symbol_ty) => match &symbol_ty.kind {
-                                        TypeKind::Union(union_types) => {
-                                            if union_types.iter().any(|ut| {
-                                                matches!(
-                                                    &ut.kind,
-                                                    TypeKind::StrLit(_) | TypeKind::Str
-                                                )
-                                            }) {
-                                                Some(KCLCompletionItemKind::Function)
-                                            } else {
-                                                type_to_item_kind(attr_ty)
-                                            }
-                                        }
                                         TypeKind::Schema(_) => {
                                             Some(KCLCompletionItemKind::SchemaAttr)
                                         }
@@ -348,7 +318,6 @@ fn completion_dot(
         }
         None => {}
     }
-
     Some(into_completion_items(&items).into())
 }