Refactor coordinate api

api/s24_coordinates.py

@@ -0,0 +1,45 @@
+from psycopg.rows import dict_row, Row
+from .connection import g_conn_dict as conn
+from .s0_base import *
+from .operation import *
+from .change_set import ChangeSet
+
+
+def _to_point(coord: str) -> dict[str, float]:
+    coord = coord.removeprefix('(')
+    coord = coord.removesuffix(')')
+    coord = coord.split(',')
+    return { 'x': float(coord[0]), 'y': float(coord[1]) }
+
+
+def get_node_coord(name: str, id: str) -> dict[str, float] | None:
+    with conn[name].cursor(row_factory=dict_row) as cur:
+        cur.execute(f"select * from coordinates where node = '{id}'")
+        row = cur.fetchone()
+        if row == None:
+            return None
+        
+        coord = str(row['coord'])
+        return _to_point(coord)
+
+
+def set_node_coord(name: str, id: str, x: float, y: float) -> ChangeSet:
+    if not is_junction(name, id):
+        return
+
+    old = get_node_coord(name, id)
+    if old == None:
+        return
+    old_x, old_y = old['x'], old['y']
+
+    with conn[name].cursor() as cur:
+        sql = f"update coordinates set coord = '({x},{y})' where node = '{id}'"
+        cur.execute(sql)
+
+        redo = sql.replace("'", '"')
+        undo = f'update coordinates set coord = "({old_x},{old_y})" where node = "{id}"'
+        add_operation(name, redo, undo)
+
+    change = ChangeSet()
+    change.update('junction', id, 'coord', 'point', str({'x': x, 'y': y}))
+    return change

api/s2_junctions.py

@@ -3,6 +3,8 @@ from .connection import g_conn_dict as conn
 from .s0_base import *
 from .operation import *
 from .change_set import ChangeSet
+from .s24_coordinates import *
+
 
 def add_junction(name: str, id: str, x: float, y: float, elevation: float) -> ChangeSet:
     if is_node(name, id):
@@ -68,10 +70,12 @@ def _get_junction(name: str, id: str) -> Row | None:
         cur.execute(f"select elevation, demand, pattern from junctions where id = '{id}'")
         return cur.fetchone()
 
+
 def get_junction_elevation(name: str, id: str) -> float | None:
     row = _get_junction(name, id)
     return float(row['elevation']) if row != None else None
 
+
 def get_junction_demand(name: str, id: str) -> float | str | None:
     row = _get_junction(name, id)
     if row != None:
@@ -79,6 +83,7 @@ def get_junction_demand(name: str, id: str) -> float | str | None:
     else:
         return None
 
+
 def get_junction_pattern(name: str, id: str) -> str | None:
     row = _get_junction(name, id)
     if row != None:
@@ -86,21 +91,10 @@ def get_junction_pattern(name: str, id: str) -> str | None:
     else:
         return None
 
-def _to_point(coord: str) -> dict[str, float]:
-    coord = coord.removeprefix('(')
-    coord = coord.removesuffix(')')
-    coord = coord.split(',')
-    return { 'x': float(coord[0]), 'y': float(coord[1]) }
 
 def get_junction_coord(name: str, id: str) -> dict[str, float] | None:
-    with conn[name].cursor(row_factory=dict_row) as cur:
+    return get_node_coord(name, id)
-        cur.execute(f"select * from coordinates where node = '{id}'")
+
-        row = cur.fetchone()
-        if row == None:
-            return None
-        
-        coord = str(row['coord'])
-        return _to_point(coord)
 
 def set_junction_elevation(name: str, id: str, elevation: float) -> ChangeSet:
     if not is_junction(name, id):
@@ -167,22 +161,4 @@ def set_junction_pattern(name: str, id: str, pattern: str) -> ChangeSet:
 
 
 def set_junction_coord(name: str, id: str, x: float, y: float) -> ChangeSet:
-    if not is_junction(name, id):
+    return set_node_coord(name, id, x, y)
-        return
-
-    old = get_junction_coord(name, id)
-    if old == None:
-        return
-    old_x, old_y = old['x'], old['y']
-
-    with conn[name].cursor() as cur:
-        sql = f"update coordinates set coord = '({x},{y})' where node = '{id}'"
-        cur.execute(sql)
-
-        redo = sql.replace("'", '"')
-        undo = f'update coordinates set coord = "({old_x},{old_y})" where node = "{id}"'
-        add_operation(name, redo, undo)
-
-    change = ChangeSet()
-    change.update('junction', id, 'coord', 'point', str({'x': x, 'y': y}))
-    return change

api/s3_reservoirs.py

@@ -3,6 +3,8 @@ from .connection import g_conn_dict as conn
 from .s0_base import *
 from .operation import *
 from .change_set import ChangeSet
+from .s24_coordinates import *
+
 
 def add_reservoir(name: str, id: str, x: float, y: float, head: float) -> ChangeSet:
     if is_node(name, id):
@@ -24,6 +26,7 @@ def add_reservoir(name: str, id: str, x: float, y: float, head: float) -> Change
     change.add('reservoir', id)
     return change
 
+
 def delete_reservoir(name: str, id: str) -> ChangeSet:
     if not is_reservoir(name, id):
         return
@@ -76,21 +79,8 @@ def get_reservoir_pattern(name: str, id: str) -> str | None:
     else:
         return None
 
-def _to_point(coord: str) -> dict[str, float]:
-    coord = coord.removeprefix('(')
-    coord = coord.removesuffix(')')
-    coord = coord.split(',')
-    return { 'x': float(coord[0]), 'y': float(coord[1]) }
-
 def get_reservoir_coord(name: str, id: str) -> dict[str, float] | None:
-    with conn[name].cursor(row_factory=dict_row) as cur:
+    return get_node_coord(name, id)
-        cur.execute(f"select * from coordinates where node = '{id}'")
-        row = cur.fetchone()
-        if row == None:
-            return None
-        
-        coord = str(row['coord'])
-        return _to_point(coord)
 
 def set_reservoir_head(name: str, id: str, head: float) -> ChangeSet:
     if not is_reservoir(name, id):
@@ -136,22 +126,4 @@ def set_reservoir_pattern(name: str, id: str, pattern: str) -> ChangeSet:
 
 
 def set_reservoir_coord(name: str, id: str, x: float, y: float) -> ChangeSet:
-    if not is_reservoir(name, id):
+    return set_node_coord(name, id, x, y)
-        return
-
-    old = get_reservoir_coord(name, id)
-    if old == None:
-        return
-    old_x, old_y = old['x'], old['y']
-
-    with conn[name].cursor() as cur:
-        sql = f"update coordinates set coord = '({x},{y})' where node = '{id}'"
-        cur.execute(sql)
-
-        redo = sql.replace("'", '"')
-        undo = f'update coordinates set coord = "({old_x},{old_y})" where node = "{id}"'
-        add_operation(name, redo, undo)
-
-    change = ChangeSet()
-    change.update('reservoir', id, 'coord', 'point', str({'x': x, 'y': y}))
-    return change

api/s4_tanks.py

@@ -3,6 +3,8 @@ from .connection import g_conn_dict as conn
 from .s0_base import *
 from .operation import *
 from .change_set import ChangeSet
+from .s24_coordinates import *
+
 
 def add_tank(name: str, id: str, x: float, y: float, elevation: float, init_level: float = 0, min_level: float = 0, max_level: float = 0, diameter: float = 0, min_vol: float = 0) -> ChangeSet:
     if is_node(name, id):
@@ -120,22 +122,8 @@ def get_tank_overflow(name: str, id: str) -> str | None:
         return None
 
 
-def _to_point(coord: str) -> dict[str, float]:
-    coord = coord.removeprefix('(')
-    coord = coord.removesuffix(')')
-    coord = coord.split(',')
-    return { 'x': float(coord[0]), 'y': float(coord[1]) }
-
-
 def get_tank_coord(name: str, id: str) -> dict[str, float] | None:
-    with conn[name].cursor(row_factory=dict_row) as cur:
+    return get_node_coord(name, id)
-        cur.execute(f"select * from coordinates where node = '{id}'")
-        row = cur.fetchone()
-        if row == None:
-            return None
-        
-        coord = str(row['coord'])
-        return _to_point(coord)
 
 
 def set_tank_elevation(name: str, id: str, elevation: float) -> ChangeSet:
@@ -312,22 +300,4 @@ def set_tank_overflow(name: str, id: str, overflow: str) -> ChangeSet:
 
 
 def set_tank_coord(name: str, id: str, x: float, y: float) -> ChangeSet:
-    if not is_tank(name, id):
+    return set_node_coord(name, id, x, y)
-        return
-
-    old = get_tank_coord(name, id)
-    if old == None:
-        return
-    old_x, old_y = old['x'], old['y']
-
-    with conn[name].cursor() as cur:
-        sql = f"update coordinates set coord = '({x},{y})' where node = '{id}'"
-        cur.execute(sql)
-
-        redo = sql.replace("'", '"')
-        undo = f'update coordinates set coord = "({old_x},{old_y})" where node = "{id}"'
-        add_operation(name, redo, undo)
-
-    change = ChangeSet()
-    change.update('tank', id, 'coord', 'point', str({'x': x, 'y': y}))
-    return change