diff --git a/api/__init__.py b/api/__init__.py index 3f756eb..b36ed17 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -29,7 +29,7 @@ from .s0_base import is_node, is_junction, is_reservoir, is_tank from .s0_base import is_link, is_pipe, is_pump, is_valve from .s0_base import is_curve from .s0_base import is_pattern -from .s0_base import get_nodes, get_junctions, get_reservoirs, get_tanks, get_links, get_pipes, get_pumps, get_valves, get_curves, get_patterns +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 .s1_title import get_title_schema, get_title, set_title diff --git a/api/s0_base.py b/api/s0_base.py index c17f9f8..52da948 100644 --- a/api/s0_base.py +++ b/api/s0_base.py @@ -98,6 +98,15 @@ def _get_nodes_by_type(name: str, type: str) -> list[str]: ids.append(record['id']) return ids +# DingZQ +def get_nodes_id_and_type(name: str) -> dict[str, str]: + nodes_id_and_type: dict[str, str] = {} + with conn[name].cursor(row_factory=dict_row) as cur: + cur.execute(f"select id, type from {_NODE} order by id") + for record in cur: + nodes_id_and_type[record['id']] = record['type'] + return nodes_id_and_type + # DingZQ def get_junctions(name: str) -> list[str]: return _get_nodes_by_type(name, JUNCTION) @@ -123,6 +132,15 @@ def _get_links_by_type(name: str, type: str) -> list[str]: ids.append(record['id']) return ids +# DingZQ +def get_links_id_and_type(name: str) -> dict[str, str]: + links_id_and_type: dict[str, str] = {} + with conn[name].cursor(row_factory=dict_row) as cur: + cur.execute(f"select id, type from {_LINK} order by id") + for record in cur: + links_id_and_type[record['id']] = record['type'] + return links_id_and_type + # DingZQ def get_pipes(name: str) -> list[str]: return _get_links_by_type(name, PIPE) diff --git a/tjnetwork.py b/tjnetwork.py index b3467af..de06930 100644 --- a/tjnetwork.py +++ b/tjnetwork.py @@ -832,23 +832,12 @@ def get_node_coord(name: str, node_id: str) -> dict[str, float]: # DingZQ, 2024-12-08, get all node coord # id, x, y, type -def get_network_coords(name: str) -> dict[str, dict[str, float]]: - node_ids = api.get_nodes(name) - junctions = api.get_junctions(name) - reservoirs = api.get_reservoirs(name) - tanks = api.get_tanks(name) +def get_network_node_coords(name: str) -> dict[str, dict[str, float]]: + nodes_id_and_type = api.get_nodes_id_and_type(name) result = {} - for node_id in node_ids: - #result[node_id] = api.get_node_coord(name, node_id) + for node_id, node_type in nodes_id_and_type.items(): coord = api.get_node_coord(name, node_id) - if node_id in junctions: - coord['type'] = 'junction' - elif node_id in reservoirs: - coord['type'] = 'reservoir' - elif node_id in tanks: - coord['type'] = 'tank' - else: - coord['type'] = 'node' + coord['type'] = node_type result[node_id] = coord return result