From 57dcabb4b023a8488174000331a0d2b575c3ebe7 Mon Sep 17 00:00:00 2001
From: Jin IGARASHI <jin_igarashi@kk-grp.jp>
Date: Mon, 28 Oct 2019 17:41:20 +0900
Subject: [PATCH] I fixed bugs to export shapefiles of valve and pumps. The
 script is going to update node id of pipeline which intersects valve and
 pump.

---
 epanet/layer_base.py | 20 +++++++++++++++++++-
 epanet/pipes.py      |  4 ++++
 epanet/pumps.py      |  1 +
 epanet/tasks.py      |  6 +++---
 epanet/valves.py     |  2 ++
 5 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/epanet/layer_base.py b/epanet/layer_base.py
index 4588769..152b9f1 100644
--- a/epanet/layer_base.py
+++ b/epanet/layer_base.py
@@ -12,4 +12,22 @@ def createProjection(self, output_dir):
         shutil.copy("./templates/wgs84_4326.prj","{0}.prj".format(output_dir))
 
     def get_file_path(self, f):
-        return "{0}/{1}".format(f.name.replace(".inp", ""), self.layer)
\ No newline at end of file
+        return "{0}/{1}".format(f.name.replace(".inp", ""), self.layer)
+
+    def updatePipeNode(self, obj):
+        '''
+        To update pipe node which intersects valve or pump node.
+        :param obj: valve object or pump object shall be here
+        :return:
+        '''
+        target_key = ",".join([str(obj.lon), str(obj.lat)])
+        if target_key in self.coords.coordMap and self.coords.coordMap[target_key]:
+            coord = self.coords.coordMap[target_key]
+            nodeid = coord.id
+            for pipe in self.pipes:
+                if nodeid == pipe.node1:
+                    pipe.set_node(obj.id, pipe.node2)
+                    coord.id = obj.id
+                elif nodeid == pipe.node2:
+                    pipe.set_node(pipe.node1, obj.id)
+                    coord.id = obj.id
\ No newline at end of file
diff --git a/epanet/pipes.py b/epanet/pipes.py
index c19b319..3a93596 100644
--- a/epanet/pipes.py
+++ b/epanet/pipes.py
@@ -16,6 +16,10 @@ def __init__(self, id, node1, node2, length, diameter):
             self.minorloss = 0
             self.status = "Open"
 
+        def set_node(self, node1, node2):
+            self.node1 = node1
+            self.node2 = node2
+
         @staticmethod
         def create_header(f):
             f.writelines("[PIPES]\n")
diff --git a/epanet/pumps.py b/epanet/pumps.py
index 0132ea8..b9aebd4 100644
--- a/epanet/pumps.py
+++ b/epanet/pumps.py
@@ -127,5 +127,6 @@ def export_shapefile(self, f):
             for p in self.pumps:
                 _shp.point(float(p.lon), float(p.lat))
                 _shp.record(p.id, p.elevation, p.curve.head, p.curve.flow, None, "POWER {0}".format(str(1)))
+                self.updatePipeNode(p)
             _shp.close()
         self.createProjection(filename)
diff --git a/epanet/tasks.py b/epanet/tasks.py
index 9846ef6..d77e5f4 100644
--- a/epanet/tasks.py
+++ b/epanet/tasks.py
@@ -86,7 +86,7 @@ def execute(self):
                 tanks.export(f)
                 #join lists of pipes id which duplicates from pumps and valves.
                 del_pipes_id = []
-                del_pipes_id.extend( pumps.get_del_pipes_id_for_inp())
+                del_pipes_id.extend(pumps.get_del_pipes_id_for_inp())
                 del_pipes_id.extend(valves.get_del_pipes_id_for_inp())
                 pipes.export(f, del_pipes_id)
                 pumps.export(f)
@@ -101,10 +101,10 @@ def execute(self):
                 del_coords_id = []
                 del_coords_id.extend(pumps.get_del_coords_id_for_inp())
                 del_coords_id.extend(valves.get_del_coords_id_for_inp())
-                coords.export_shapefile(f, del_coords_id)
-                pipes.export_shapefile(f)
                 tanks.export_shapefile(f)
                 reservoirs.export_shapefile(f)
                 pumps.export_shapefile(f)
                 valves.export_shapefile(f)
+                coords.export_shapefile(f, del_coords_id)
+                pipes.export_shapefile(f)
                 shutil.copy("./templates/template_qgs_project.qgz", "{0}/{1}_{2}.qgz".format(f.name.replace(".inp", ""), self.wss_id, self.wss_name))
diff --git a/epanet/valves.py b/epanet/valves.py
index 13bada4..b296b0f 100644
--- a/epanet/valves.py
+++ b/epanet/valves.py
@@ -110,5 +110,7 @@ def export_shapefile(self, f):
             for v in self.valves:
                 _shp.point(float(v.lon), float(v.lat))
                 _shp.record(v.id, v.elevation, v.diameter, v.valve_type, v.setting, v.minor_loss)
+                self.updatePipeNode(v)
             _shp.close()
         self.createProjection(filename)
+