From f8ac46d80d9ab5d06aae63e0830604c584b62cde Mon Sep 17 00:00:00 2001
From: DingZQ <ding@aecforge.com>
Date: Wed, 18 Dec 2024 22:28:57 +0800
Subject: [PATCH] Add methods to get node types

---
 api/s0_base.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/api/s0_base.py b/api/s0_base.py
index 6612555..4dab020 100644
--- a/api/s0_base.py
+++ b/api/s0_base.py
@@ -89,11 +89,54 @@ def _get_all(name: str, base_type: str) -> list[str]:
 def get_nodes(name: str) -> list[str]:
     return _get_all(name, _NODE)
 
+# DingZQ
+def _get_nodes_by_type(name: str, type: str) -> list[str]:
+    ids : list[str] = []
+    with conn[name].cursor(row_factory=dict_row) as cur:
+        cur.execute(f"select id from {_NODE} where type = '{type}' order by id")
+        for record in cur:
+            ids.append(record['id'])
+    return ids 
 
+# DingZQ
+def get_junctions(name: str) -> list[str]:
+    return _get_nodes_by_type(name, JUNCTION)
+
+# DingZQ
+def get_reservoir(name: str) -> list[str]:
+    return _get_nodes_by_type(name, RESERVOIR)
+
+# DingZQ
+def get_tank(name: str) -> list[str]:
+    return _get_nodes_by_type(name, TANK)
+
+# DingZQ
 def get_links(name: str) -> list[str]:
     return _get_all(name, _LINK)
 
 
+# DingZQ
+def _get_links_by_type(name: str, type: str) -> list[str]:
+    ids : list[str] = []
+    with conn[name].cursor(row_factory=dict_row) as cur:
+        cur.execute(f"select id from {_LINK} where type = '{type}' order by id")
+        for record in cur:
+            ids.append(record['id'])
+    return ids
+
+# DingZQ
+def get_pipes(name: str) -> list[str]:
+    return _get_links_by_type(name, PIPE)
+
+# DingZQ
+def get_pumps(name: str) -> list[str]:
+    return _get_links_by_type(name, PUMP)
+
+# DingZQ
+def get_valves(name: str) -> list[str]:
+    return _get_links_by_type(name, VALVE)
+
+
 def get_curves(name: str) -> list[str]:
     return _get_all(name, _CURVE)