Merge pull request #186 from BuczynskiRafal/read_tags

Add getter and setter for tags.

swmmio/core.py

@@ -188,6 +188,7 @@ def conduits(self):
         # create dataframes of relevant sections from the INP
         conduits_df = dataframe_from_inp(inp.path, "CONDUITS")
         xsections_df = dataframe_from_inp(inp.path, "XSECTIONS")
+        tags = dataframe_from_inp(inp.path, "TAGS")
         conduits_df = conduits_df.join(xsections_df)
 
         if rpt:
@@ -212,6 +213,10 @@ def conduits(self):
         df.InletNode = df.InletNode.astype(str)
         df.OutletNode = df.OutletNode.astype(str)
 
+        tags = dataframe_from_inp(inp.path, "TAGS")
+        if "Link" in set(tags.index):
+            df = df.merge(dataframe_from_inp(inp.path, "Tags"), left_on="Name", right_on="Name", how="left")
+
         self._conduits_df = df
 
         return df
@@ -344,6 +349,7 @@ def subcatchments(self):
             return self._subcatchments_df
 
         df = ModelSection(model=self, **COMPOSITE_OBJECTS['subcatchments'])
+
         self._subcatchments_df = df
         return df
 
@@ -534,6 +540,7 @@ def __init__(self, file_path):
         self._inflows_df = None
         self._curves_df = None
         self._timeseries_df = None
+        self._tags_df = None
 
         SWMMIOFile.__init__(self, file_path)  # run the superclass init
 
@@ -568,7 +575,8 @@ def __init__(self, file_path):
             '[HYDROGRAPHS]',
             '[INFLOWS]',
             '[Polygons]',
-            '[TIMESERIES]'
+            '[TIMESERIES]',
+            '[TAGS]',
         ]
 
     def save(self, target_path=None):
@@ -1314,6 +1322,20 @@ def timeseries(self, df):
         """Set inp.timeseries DataFrame."""
         self._timeseries_df = df
 
+    @property
+    def tags(self):
+        """
+        Get/set tags section of the INP file.
+        """
+        if self._tags_df is None:
+            self._tags_df = dataframe_from_inp(self.path, "[TAGS]")
+        return self._tags_df
+
+    @tags.setter
+    def tags(self, df):
+        """Set inp.tags DataFrame."""
+        self._tags_df = df
+
 
 def drop_invalid_model_elements(inp):
     """

swmmio/defs/section_headers.yml

@@ -176,5 +176,7 @@ composite:
     join_sections: [subareas]
     rpt_sections: [Subcatchment Runoff Summary]
     geomtype: polygon
-
+  tags:
+    inp_sections: [tags]
+    columns: Name, Tag
 

swmmio/examples.py

@@ -1,10 +1,11 @@
 from swmmio import Model
 from swmmio.tests.data import (MODEL_A_PATH, MODEL_EX_1, MODEL_FULL_FEATURES_XY,
-                               MODEL_FULL_FEATURES__NET_PATH, MODEL_FULL_FEATURES_XY_B)
+                               MODEL_FULL_FEATURES__NET_PATH, MODEL_FULL_FEATURES_XY_B, MODEL_GREEN_AMPT)
 
 # example models
 philly = Model(MODEL_A_PATH, crs="+init=EPSG:2817")
 jersey = Model(MODEL_FULL_FEATURES_XY)
 jerzey = Model(MODEL_FULL_FEATURES_XY_B)
 spruce = Model(MODEL_FULL_FEATURES__NET_PATH)
 walnut = Model(MODEL_EX_1)
+green = Model(MODEL_GREEN_AMPT)

swmmio/tests/test_model_elements.py

@@ -6,7 +6,7 @@
 import swmmio
 import swmmio.utils.functions
 import swmmio.utils.text
-from swmmio.tests.data import MODEL_FULL_FEATURES_XY, MODEL_FULL_FEATURES__NET_PATH, MODEL_A_PATH, MODEL_EX_1, MODEL_EX_1B
+from swmmio.tests.data import MODEL_FULL_FEATURES_XY, MODEL_FULL_FEATURES__NET_PATH, MODEL_A_PATH, MODEL_EX_1, MODEL_EX_1B, MODEL_GREEN_AMPT
 from swmmio import Model, dataframe_from_inp
 
 from swmmio.utils.dataframes import get_link_coords
@@ -171,4 +171,31 @@ def test_get_set_timeseries(test_model_02):
     assert(ts.loc['TS2'].Date == 'FILE')
     assert('"' in ts.loc['TS2'].Value)
     assert(ts.Value.isnull().sum() == 0)
-    print (ts)
+    print (ts)
+
+
+def test_inp_tags_getter_and_setter():
+    common_data = {'ElementType': ['Subcatch'] * 4,
+                   'Name': ['CA-1', 'CA-7', 'CA-8', 'CA-11'],
+                   'Tag': ['CA'] * 4,
+                   }
+
+    expected_output = pd.DataFrame(common_data)
+    expected_output.set_index(['ElementType'], inplace=True)
+
+    common_data['Tag'] = ['Modified'] * 4
+    tags_to_set = pd.DataFrame(common_data)
+    tags_to_set.set_index(['ElementType'], inplace=True)
+
+    model = Model(MODEL_GREEN_AMPT)
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        temp_inp_path = os.path.join(tempdir, f'{model.inp.name}.inp')
+        model.inp.save(temp_inp_path)
+        temp_model = Model(temp_inp_path)
+
+        assert expected_output.equals(temp_model.inp.tags.sort_index())
+
+        temp_model.inp.tags["Tag"] = ["Modified"] * 4
+
+        assert tags_to_set.equals(temp_model.inp.tags.sort_index())