Merge pull request #138 from aerispaha/bugfix/134

update pyshp API refernce in export_to_shapefile

swmmio/tests/test_spatial.py

@@ -0,0 +1,20 @@
+import unittest
+import tempfile
+import os
+
+from swmmio.examples import philly
+
+
+class TestSpatialFunctions(unittest.TestCase):
+    def setUp(self) -> None:
+        self.test_dir = tempfile.gettempdir()
+
+    def test_write_shapefile(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+
+            philly.export_to_shapefile(tmp_dir)
+            nodes_path = os.path.join(tmp_dir, f'{philly.name}_nodes.shp')
+            links_path = os.path.join(tmp_dir, f'{philly.name}_conduits.shp')
+
+            self.assertTrue(os.path.exists(nodes_path))
+            self.assertTrue(os.path.exists(links_path))

swmmio/utils/spatial.py

@@ -200,25 +200,24 @@ def write_shapefile(df, filename, geomtype='line', prj=None):
 
     # create a shp file writer object of geom type 'point'
     if geomtype == 'point':
-        w = shapefile.Writer(shapefile.POINT)
+        w = shapefile.Writer(filename, shapefile.POINT, autoBalance=True)
     elif geomtype == 'line':
-        w = shapefile.Writer(shapefile.POLYLINE)
+        w = shapefile.Writer(filename, shapefile.POLYLINE, autoBalance=True)
     elif geomtype == 'polygon':
-        w = shapefile.Writer(shapefile.POLYGON)
-
-    # use the helper mode to ensure the # of records equals the # of shapes
-    # (shapefile are made up of shapes and records, and need both to be valid)
-    w.autoBalance = 1
+        w = shapefile.Writer(filename, shapefile.POLYGON, autoBalance=True)
 
     # add the fields
     for fieldname in df.columns:
         w.field(fieldname, "C")
 
     for k, row in df.iterrows():
         w.record(*row.tolist())
-        w.line(parts=[row.coords])
+        if geomtype == 'line':
+            w.line([row.coords])
+        if geomtype == 'point':
+            w.point(*row.coords[0])
 
-    w.save(filename)
+    w.close()
 
     # add projection data to the shapefile,
     if prj is None: