update: using PkgType in GraphNode

TEN-framework · Sep 23, 2024 · 641cae7 · 641cae7
1 parent 14777d0
commit 641cae7
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 16 deletions.
diff --git a/core/src/ten_rust/src/pkg_info/graph/check/connections_are_compatible.rs b/core/src/ten_rust/src/pkg_info/graph/check/connections_are_compatible.rs
@@ -28,7 +28,7 @@ impl Graph {
         self.nodes
             .iter()
             .find_map(|node| {
-                if node.node_type == "extension"
+                if node.node_type == PkgType::Extension
                     && node.name.as_str() == extension
                     && node.app.as_str() == app
                 {

diff --git a/core/src/ten_rust/src/pkg_info/graph/check/connections_nodes_are_defined.rs b/core/src/ten_rust/src/pkg_info/graph/check/connections_nodes_are_defined.rs
@@ -6,7 +6,10 @@
 //
 use anyhow::Result;
 
-use crate::pkg_info::graph::{Graph, GraphMessageFlow};
+use crate::pkg_info::{
+    graph::{Graph, GraphMessageFlow},
+    pkg_type::PkgType,
+};
 
 impl Graph {
     fn check_if_dest_of_connection_are_defined_in_nodes(
@@ -47,7 +50,7 @@ impl Graph {
 
         let mut all_extensions: Vec<String> = Vec::new();
         for node in &self.nodes {
-            if node.node_type.as_str() == "extension" {
+            if node.node_type == PkgType::Extension {
                 let unique_ext_name = format!(
                     "{}:{}:{}",
                     node.app.as_str(),

diff --git a/core/src/ten_rust/src/pkg_info/graph/check/duplicated_nodes.rs b/core/src/ten_rust/src/pkg_info/graph/check/duplicated_nodes.rs
@@ -6,7 +6,7 @@
 //
 use anyhow::Result;
 
-use crate::pkg_info::graph::Graph;
+use crate::pkg_info::{graph::Graph, pkg_type::PkgType};
 
 impl Graph {
     pub fn check_if_nodes_duplicated(&self) -> Result<()> {
@@ -15,8 +15,8 @@ impl Graph {
         let mut all_extension_groups: Vec<String> = Vec::new();
 
         for (node_idx, node) in self.nodes.iter().enumerate() {
-            match node.node_type.as_str() {
-                "extension" => {
+            match node.node_type {
+                PkgType::Extension => {
                     let unique_ext_name = format!(
                         "{}:{}:{}",
                         node.app.as_str(),
@@ -44,7 +44,7 @@ impl Graph {
                     }
                 }
 
-                "extension_group" => {
+                PkgType::ExtensionGroup => {
                     let unique_ext_group_name =
                         format!("{}:{}", node.app.as_str(), node.name);
 

diff --git a/core/src/ten_rust/src/pkg_info/graph/check/nodes_are_installed.rs b/core/src/ten_rust/src/pkg_info/graph/check/nodes_are_installed.rs
@@ -8,15 +8,15 @@ use std::collections::HashMap;
 
 use anyhow::Result;
 
-use crate::pkg_info::{graph::Graph, PkgInfo};
+use crate::pkg_info::{graph::Graph, pkg_type::PkgType, PkgInfo};
 
 impl Graph {
     pub fn check_if_nodes_have_installed(
         &self,
         all_needed_pkgs: &HashMap<String, Vec<PkgInfo>>,
     ) -> Result<()> {
         // app, node_type, node_addon
-        let mut not_installed_pkgs: Vec<(String, String, String)> = Vec::new();
+        let mut not_installed_pkgs: Vec<(String, PkgType, String)> = Vec::new();
 
         for node in &self.nodes {
             if !all_needed_pkgs.contains_key(node.app.as_str()) {
@@ -29,7 +29,7 @@ impl Graph {
 
             let pkgs_in_app = all_needed_pkgs.get(node.app.as_str()).unwrap();
             let found = pkgs_in_app.iter().find(|pkg| {
-                pkg.pkg_identity.pkg_type.to_string() == node.node_type
+                pkg.pkg_identity.pkg_type == node.node_type
                     && pkg.pkg_identity.name == node.addon
                     && pkg.is_local_installed
             });

diff --git a/core/src/ten_rust/src/pkg_info/graph/mod.rs b/core/src/ten_rust/src/pkg_info/graph/mod.rs
@@ -12,6 +12,7 @@ use anyhow::Result;
 use serde::{Deserialize, Serialize};
 
 use super::{
+    pkg_type::PkgType,
     predefined_graphs::{
         connection::{PkgConnection, PkgDestination, PkgMessageFlow},
         node::PkgNode,
@@ -76,7 +77,7 @@ impl Graph {
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct GraphNode {
     #[serde(rename = "type")]
-    pub node_type: String,
+    pub node_type: PkgType,
     pub name: String,
     pub addon: String,
 
@@ -95,7 +96,9 @@ pub struct GraphNode {
 impl GraphNode {
     fn validate_and_complete(&mut self) -> Result<()> {
         // extension node must specify extension_group name.
-        if self.node_type == "extension" && self.extension_group.is_none() {
+        if self.node_type == PkgType::Extension
+            && self.extension_group.is_none()
+        {
             return Err(anyhow::anyhow!(
                 "Node '{}' of type 'extension' must have an 'extension_group' defined.",
                 self.name
@@ -109,7 +112,7 @@ impl GraphNode {
 impl From<PkgNode> for GraphNode {
     fn from(pkg_node: PkgNode) -> Self {
         GraphNode {
-            node_type: pkg_node.node_type.to_string(),
+            node_type: pkg_node.node_type.clone(),
             name: pkg_node.name.clone(),
             addon: pkg_node.addon.clone(),
             extension_group: pkg_node.extension_group.clone(),

diff --git a/core/src/ten_rust/src/pkg_info/predefined_graphs/node.rs b/core/src/ten_rust/src/pkg_info/predefined_graphs/node.rs
@@ -4,8 +4,6 @@
 // Licensed under the Apache License, Version 2.0, with certain conditions.
 // Refer to the "LICENSE" file in the root directory for more information.
 //
-use std::str::FromStr;
-
 use crate::pkg_info::{graph::GraphNode, pkg_type::PkgType, PkgInfo};
 
 #[derive(Debug, Clone)]
@@ -28,7 +26,7 @@ pub struct PkgNode {
 impl From<GraphNode> for PkgNode {
     fn from(manifest_node: GraphNode) -> Self {
         PkgNode {
-            node_type: PkgType::from_str(&manifest_node.node_type).unwrap(),
+            node_type: manifest_node.node_type,
             name: manifest_node.name,
             addon: manifest_node.addon,
             extension_group: manifest_node.extension_group,