organized tree code

Author: nicoschmdt

core/libs/commonwealth/src/commonwealth/utils/tree.py

@@ -1,39 +1,21 @@
-from typing import Any, Callable, List, Dict
+from typing import Any, Callable, Dict, List
+
+import zenoh
 
 
 class TreeNode:
     def __init__(self, segment: str):
         self.segment = segment
-        self.children: Dict[str, 'TreeNode'] = {}
+        self.children: Dict[str, "TreeNode"] = {}
         self.is_valid = False
         self.methods: Dict[str, Callable[..., Any]] = {}
-    
-    def __str__(self, level=0, prefix="", is_last=True):
-        ret = prefix
-        if level > 0:
-            ret += "|___ " if is_last else "|--- "
-        ret += f"{self.segment}"
-        if self.is_valid:
-            ret += " [Valid]"
-        ret += "\n"
-        
-        prefix += "    " if is_last else "|   "
-        children_list = list(self.children.values())
-        for i, child in enumerate(children_list):
-            is_last_child = (i == len(children_list) - 1)
-            ret += child.__str__(level + 1, prefix, is_last_child)
-        
-        return ret
-
-    def add_child(self, child: 'TreeNode') -> 'TreeNode':
+
+    def add_child(self, child: "TreeNode") -> "TreeNode":
         if child.segment in self.children:
             return self.children[child.segment]
         self.children[child.segment] = child
         return child
-    
-    def is_valid(self) -> bool:
-        return self.is_valid
-    
+
     def get_methods(self) -> Dict[str, Callable[..., Any]]:
         return self.methods
 
@@ -45,20 +27,20 @@ def add_node(self, segments: List[str], method: str, func: Callable[..., Any]) -
             self.is_valid = True
             self.methods[method] = func
             return
-            
+
         child = self.add_child(TreeNode(segments[1]))
         child.add_node(segments[1:], method, func)
-        
+
     def process_path(self, path: str, method: str, func: Callable[..., Any]) -> None:
-        segments = path.split('/')
+        segments = path.split("/")
         self.add_node(segments, method, func)
 
-    def get_match(self, path: str) -> tuple[str, Dict[str, Callable[..., Any]]] | None:
-        segments = path.split('/')
+    def get_match(self, path: str) -> tuple[str, "TreeNode"] | None:
+        segments = path.split("/")
 
         if self.segment != segments[0]:
             return None
-        
+
         node = self
         matched_path = node.segment
 
@@ -72,8 +54,19 @@ def get_match(self, path: str) -> tuple[str, Dict[str, Callable[..., Any]]] | No
                     matched_path += "/" + node.segment
                 except KeyError:
                     return None
-        
 
         if node.is_valid:
-            return matched_path, node.get_methods()
+            return matched_path, node
         return None
+
+    def get_corresponding_method(self, sampleKind: zenoh.SampleKind) -> Callable[..., Any] | None:
+        if sampleKind == zenoh.SampleKind.DELETE:
+            keyword = "DELETE"
+        else:
+            methods = [method for method, _ in self.methods.items() if method != "DELETE"]
+            keyword = methods[0]
+
+        try:
+            return self.methods[keyword]
+        except KeyError:
+            return None