Add API to get major nodes and major pipes

api/__init__.py

@@ -32,6 +32,7 @@ from .s0_base import is_curve
 from .s0_base import is_pattern
 from .s0_base import get_nodes, get_nodes_id_and_type, get_junctions, get_reservoirs, get_tanks, get_links, get_links_id_and_type, get_pipes, get_pumps, get_valves, get_curves, get_patterns
 from .s0_base import get_node_links, get_link_nodes
+from .s0_base import get_major_nodes, get_major_pipes
 
 from .s1_title import get_title_schema, get_title, set_title
 

api/s0_base.py

@@ -107,7 +107,18 @@ def get_nodes_id_and_type(name: str) -> dict[str, str]:
             nodes_id_and_type[record['id']] = record['type']
     return nodes_id_and_type
 
-# DingZQ
+# DingZQ 2024-12-31
+def get_major_nodes(name: str) -> list[str]:
+    major_nodes_set = set()
+    with conn[name].cursor(row_factory=dict_row) as cur:
+        cur.execute(f"select node1, node2 from pipes where diameter > 800")
+        for record in cur:
+            major_nodes_set.add(record['node1'])
+            major_nodes_set.add(record['node2'])
+
+    return list(major_nodes_set)
+
+# DingZQs
 def get_junctions(name: str) -> list[str]:
     return _get_nodes_by_type(name, JUNCTION)
 
@@ -141,6 +152,16 @@ def get_links_id_and_type(name: str) -> dict[str, str]:
             links_id_and_type[record['id']] = record['type']
     return links_id_and_type
 
+# DingZQ 2024-12-31
+# 获取直径大于800的管道
+def get_major_pipes(name: str) -> list[str]:
+    major_pipe_ids: list[str] = []
+    with conn[name].cursor(row_factory=dict_row) as cur:
+        cur.execute(f"select id from pipes where diameter > 800 order by id")
+        for record in cur:
+            major_pipe_ids.append(record['id'])
+    return major_pipe_ids
+    
 # DingZQ
 def get_pipes(name: str) -> list[str]:
     return _get_links_by_type(name, PIPE)