From 6e5e24c527e0739d80ad237e576773c49b5d1538 Mon Sep 17 00:00:00 2001
From: Mike Taves <mwtoews@gmail.com>
Date: Thu, 15 Oct 2020 17:02:45 +1300
Subject: [PATCH] refactor(shapefile): add check_version for pyshp>=2; skip
 tests

---
 autotest/t006_test.py           | 16 +++++++---
 autotest/t007_test.py           | 54 ++++++++++++++++++++++-----------
 autotest/t009_test.py           |  9 ++++++
 autotest/t031_test.py           | 13 +++++++-
 autotest/t061_test_gridgen.py   | 10 ++++++
 autotest/t504_test.py           | 10 +++++-
 autotest/t505_test.py           | 11 +++++--
 autotest/t550_test.py           | 20 +++++++++---
 flopy/export/shapefile_utils.py | 50 +++++++++++++++++++++---------
 flopy/utils/geometry.py         |  2 +-
 flopy/utils/gridgen.py          |  7 ++---
 11 files changed, 152 insertions(+), 50 deletions(-)

diff --git a/autotest/t006_test.py b/autotest/t006_test.py
index aab2555ec8..1b2dfdd5ec 100644
--- a/autotest/t006_test.py
+++ b/autotest/t006_test.py
@@ -7,7 +7,12 @@
     matplotlib = None
 
 import flopy
-import shapefile
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
 
 
 cpth = os.path.join('temp', 't006')
@@ -82,10 +87,11 @@ def test_mflist_reference():
     ghb = flopy.modflow.ModflowGhb(ml, stress_period_data=ghb_dict)
     assert isinstance(ghb, flopy.modflow.ModflowGhb)
 
-    test = os.path.join(cpth, 'test3.shp')
-    ml.export(test, kper=0)
-    shp = shapefile.Reader(test)
-    assert shp.numRecords == nrow * ncol
+    if shapefile:
+        test = os.path.join(cpth, 'test3.shp')
+        ml.export(test, kper=0)
+        shp = shapefile.Reader(test)
+        assert shp.numRecords == nrow * ncol
 
 def test_cbc_ts():
     fpth = os.path.join('..', 'examples', 'data', 'mf2005_test',
diff --git a/autotest/t007_test.py b/autotest/t007_test.py
index 5ef574b3e0..e8e8be2bd2 100644
--- a/autotest/t007_test.py
+++ b/autotest/t007_test.py
@@ -33,6 +33,16 @@
     os.makedirs(spth)
 
 
+def import_shapefile():
+    try:
+        import shapefile
+    except ImportError:
+        return None
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        return None
+    return shapefile
+
+
 def remove_shp(shpname):
     os.remove(shpname)
     for ext in ['prj', 'shx', 'dbf']:
@@ -94,9 +104,8 @@ def export_netcdf(m):
 
 def export_shapefile(namfile):
     print('in export_shapefile: {}'.format(namfile))
-    try:
-        import shapefile as shp
-    except:
+    shp = import_shapefile()
+    if shp is None:
         return
 
     m = flopy.modflow.Modflow.load(namfile, model_ws=pth, verbose=False)
@@ -128,9 +137,8 @@ def export_shapefile(namfile):
 
 def export_shapefile_modelgrid_override(namfile):
     print('in export_modelgrid_override: {}'.format(namfile))
-    try:
-        import shapefile as shp
-    except:
+    shp = import_shapefile()
+    if shp is None:
         return
 
     from flopy.discretization import StructuredGrid
@@ -164,6 +172,8 @@ def export_shapefile_modelgrid_override(namfile):
 
 
 def test_output_helper_shapefile_export():
+    if import_shapefile() is None:
+        return
     ws = os.path.join('..', 'examples', 'data', 'freyberg_multilayer_transient')
     name = 'freyberg.nam'
 
@@ -177,6 +187,9 @@ def test_output_helper_shapefile_export():
 
 
 def test_freyberg_export():
+    if import_shapefile() is None:
+        return
+
     from flopy.discretization import StructuredGrid
     namfile = 'freyberg.nam'
 
@@ -282,6 +295,10 @@ def test_export_output():
 
 
 def test_write_shapefile():
+    sf = import_shapefile()
+    if not sf:
+        return
+
     from flopy.discretization import StructuredGrid
     from flopy.export.shapefile_utils import shp2recarray
     from flopy.export.shapefile_utils import write_grid_shapefile
@@ -298,7 +315,6 @@ def test_write_shapefile():
     # check that pyshp reads integers
     # this only check that row/column were recorded as "N"
     # not how they will be cast by python or numpy
-    import shapefile as sf
     sfobj = sf.Reader(outshp)
     for f in sfobj.fields:
         if f[0] == 'row' or f[0] == 'column':
@@ -323,11 +339,8 @@ def test_write_shapefile():
 
 
 def test_shapefile_polygon_closed():
-    import os
-    import flopy
-    try:
-        import shapefile
-    except:
+    shapefile = import_shapefile()
+    if shapefile is None:
         return
 
     xll, yll = 468970, 3478635
@@ -372,7 +385,9 @@ def test_export_array():
                        m.dis.top.array, nodata=nodata)
     arr = np.loadtxt(os.path.join(tpth, 'fb.asc'), skiprows=6)
 
-    m.modelgrid.write_shapefile(os.path.join(tpth, 'grid.shp'))
+    if import_shapefile() is not None:
+        m.modelgrid.write_shapefile(os.path.join(tpth, 'grid.shp'))
+
     # check bounds
     with open(os.path.join(tpth, 'fb.asc')) as src:
         for line in src:
@@ -1504,12 +1519,11 @@ def test_wkt_parse():
 
 
 def test_shapefile_ibound():
+    shapefile = import_shapefile()
+    if not shapefile:
+        return
     import os
     import flopy
-    try:
-        import shapefile
-    except:
-        return
 
     shape_name = os.path.join(spth, "test.shp")
     nam_file = "freyberg.nam"
@@ -1556,6 +1570,8 @@ def build_sfr_netcdf():
 
 
 def test_export_array2():
+    if import_shapefile() is None:
+        return
     from flopy.discretization import StructuredGrid
     from flopy.export.utils import export_array
     nrow = 7
@@ -1589,6 +1605,8 @@ def test_export_array2():
 
 
 def test_export_array_contours():
+    if import_shapefile() is None:
+        return
     from flopy.discretization import StructuredGrid
     from flopy.export.utils import export_array_contours
     nrow = 7
@@ -1626,6 +1644,8 @@ def test_export_contourf():
         import shapely
     except:
         return
+    if import_shapefile() is None:
+        return
     import matplotlib.pyplot as plt
     from flopy.export.utils import export_contourf
     filename = os.path.join(spth, 'myfilledcontours.shp')
diff --git a/autotest/t009_test.py b/autotest/t009_test.py
index f242fde553..ab13fc3427 100644
--- a/autotest/t009_test.py
+++ b/autotest/t009_test.py
@@ -9,6 +9,13 @@
 import numpy as np
 from flopy.utils.recarray_utils import create_empty_recarray
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
+
 try:
     import matplotlib
     # if os.getenv('TRAVIS'):  # are we running https://travis-ci.org/ automated tests ?
@@ -290,6 +297,8 @@ def test_export():
                                  delr=m.dis.delr.array,
                                  xoff=sr.xll, yoff=sr.yll)
     # m.sr.origin_loc = "ll"
+    if not shapefile:
+        return  # skip
     m.export(os.path.join(outpath, 'grid.shp'))
     r, d = create_sfr_data()
     sfr = flopy.modflow.ModflowSfr2(m, reach_data=r, segment_data={0: d})
diff --git a/autotest/t031_test.py b/autotest/t031_test.py
index aa5b453c27..16d864529b 100644
--- a/autotest/t031_test.py
+++ b/autotest/t031_test.py
@@ -15,6 +15,13 @@
 from flopy.utils.reference import SpatialReference
 from flopy.modpath.mp6sim import StartingLocationsFile
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
+
 mffiles = glob.glob('../examples/data/mp6/EXAMPLE*')
 path = os.path.join('temp', 't031')
 
@@ -103,7 +110,8 @@ def test_get_destination_data():
 
     # test deprecation
     sr2 = SpatialReference(xll=mg.xoffset, yll=mg.yoffset, rotation=-30)
-    m.dis.export(path + '/dis.shp')
+    if shapefile:
+        m.dis.export(path + '/dis.shp')
 
     pthld = PathlineFile(os.path.join(path, 'EXAMPLE-3.pathline'))
     epd = EndpointFile(os.path.join(path, 'EXAMPLE-3.endpoint'))
@@ -124,6 +132,9 @@ def test_get_destination_data():
                              dtype=starting_locs.dtype)
     assert np.all(np.in1d(starting_locs, pathline_locs))
 
+    if shapefile is None:
+        return  # skip remainder
+
     # test writing a shapefile of endpoints
     epd.write_shapefile(well_epd, direction='starting',
                         shpname=os.path.join(path, 'starting_locs.shp'),
diff --git a/autotest/t061_test_gridgen.py b/autotest/t061_test_gridgen.py
index 2d067fb74c..d1289b2de6 100644
--- a/autotest/t061_test_gridgen.py
+++ b/autotest/t061_test_gridgen.py
@@ -4,6 +4,13 @@
 import flopy
 from flopy.utils.gridgen import Gridgen
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
+
 cpth = os.path.join('temp', 't061')
 # delete the directory if it exists
 if os.path.isdir(cpth):
@@ -57,6 +64,9 @@ def test_gridgen():
     gu = Gridgen(dis_usg, model_ws=gridgen_ws, exe_name=exe_name,
                  vertical_pass_through=True)
 
+    if shapefile is None:
+        return  # skip remainder
+
     rf0shp = os.path.join(gridgen_ws, 'rf0')
     xmin = 7 * delr
     xmax = 12 * delr
diff --git a/autotest/t504_test.py b/autotest/t504_test.py
index 5f2448845a..d55c4eee95 100644
--- a/autotest/t504_test.py
+++ b/autotest/t504_test.py
@@ -8,6 +8,13 @@
 from flopy.mf6.modflow.mfsimulation import MFSimulation
 from flopy.mf6.mfbase import VerbosityLevel
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
+
 try:
     import pymake
 except:
@@ -179,7 +186,8 @@ def test003_gwfs_disv():
         assert array_util.array_comp(budget_fjf_valid, budget_frf)
 
     model = sim.get_model(model_name)
-    model.export('{}/{}.shp'.format(pth, test_ex_name))
+    if shapefile:
+        model.export('{}/{}.shp'.format(pth, test_ex_name))
 
     # change some settings
     chd_head_left = model.get_package('CHD_LEFT')
diff --git a/autotest/t505_test.py b/autotest/t505_test.py
index e81c5793dc..14888d3043 100644
--- a/autotest/t505_test.py
+++ b/autotest/t505_test.py
@@ -35,6 +35,12 @@
 from flopy.mf6.utils import testutils
 from flopy.mf6.mfbase import MFDataException
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
 
 try:
     import pymake
@@ -2110,8 +2116,9 @@ def test006_gwf3_disv():
 
         # export to netcdf - temporarily disabled
         # model.export(os.path.join(run_folder, "test006_gwf3.nc"))
-        # export to shape file
-        model.export(os.path.join(run_folder, "test006_gwf3.shp"))
+        if shapefile:
+            # export to shape file
+            model.export(os.path.join(run_folder, "test006_gwf3.shp"))
 
         # clean up
         sim.delete_output_files()
diff --git a/autotest/t550_test.py b/autotest/t550_test.py
index 9af272f780..c05a3a498b 100644
--- a/autotest/t550_test.py
+++ b/autotest/t550_test.py
@@ -7,6 +7,13 @@
 from flopy.utils import SpatialReference as OGsr
 from flopy.export.shapefile_utils import shp2recarray
 
+try:
+    import shapefile
+    if int(shapefile.__version__.split('.')[0]) < 2:
+        shapefile = None
+except ImportError:
+    shapefile = None
+
 tmpdir = 'temp/t550/'
 if not os.path.isdir(tmpdir):
     os.makedirs(tmpdir)
@@ -79,9 +86,10 @@ def cellid(k, i, j, nrow, ncol):
     #irch = np.zeros((nrow, ncol))
     riv6 = fp6.ModflowGwfriv(gwf, stress_period_data=spd6)
     rch6 = fp6.ModflowGwfrcha(gwf, recharge=rech)
-    #rch6.export('{}/mf6.shp'.format(tmpdir))
-    m.export('{}/mfnwt.shp'.format(tmpdir))
-    gwf.export('{}/mf6.shp'.format(tmpdir))
+    if shapefile:
+        #rch6.export('{}/mf6.shp'.format(tmpdir))
+        m.export('{}/mfnwt.shp'.format(tmpdir))
+        gwf.export('{}/mf6.shp'.format(tmpdir))
 
     riv6spdarrays = dict(riv6.stress_period_data.masked_4D_arrays_itr())
     rivspdarrays = dict(riv.stress_period_data.masked_4D_arrays_itr())
@@ -89,6 +97,9 @@ def cellid(k, i, j, nrow, ncol):
         assert np.abs(np.nansum(v) - np.nansum(riv6spdarrays[k])) < 1e-6, "variable {} is not equal".format(k)
         pass
 
+    if shapefile is None:
+        return  # skip remainder
+
     # check that the two shapefiles are the same
     ra = shp2recarray('{}/mfnwt.shp'.format(tmpdir))
     ra6 = shp2recarray('{}/mf6.shp'.format(tmpdir))
@@ -133,7 +144,8 @@ def test_huge_shapefile():
                         nper=nper, perlen=perlen, nstp=nstp,
                         tsmult=tsmult,
                         top=top, botm=botm)
-    m.export('{}/huge.shp'.format(tmpdir))
+    if shapefile:
+        m.export('{}/huge.shp'.format(tmpdir))
 
 
 if __name__ == '__main__':
diff --git a/flopy/export/shapefile_utils.py b/flopy/export/shapefile_utils.py
index c1998dcdf5..b0dd7fc264 100755
--- a/flopy/export/shapefile_utils.py
+++ b/flopy/export/shapefile_utils.py
@@ -3,6 +3,7 @@
 """
 import copy
 import shutil
+import inspect
 import json
 import numpy as np
 import os
@@ -16,16 +17,40 @@
 srefhttp = "https://spatialreference.org"
 
 
-def import_shapefile():
-    try:
-        import shapefile as sf
+def import_shapefile(check_version=True):
+    """Import shapefile module from pyshp.
 
-        return sf
-    except Exception as e:
-        raise Exception(
-            "io.to_shapefile(): error "
-            + "importing shapefile - try pip install pyshp"
+    Parameters
+    ----------
+    check_version : bool
+        Checks to ensure that pyshp is at least version 2. Default True,
+        which is usually required for Writer (which has a different API), but
+        can be False if only using Reader.
+
+    Returns
+    -------
+    module
+
+    Raises
+    ------
+    ImportError
+        If shapefile module is not found, or major version is less than 2.
+    """
+    try:
+        import shapefile
+    except ImportError:
+        raise ImportError(
+            inspect.getouterframes(inspect.currentframe())[1][3]
+            + ": error importing shapefile; try pip install pyshp"
         )
+    if check_version:
+        if int(shapefile.__version__.split(".")[0]) < 2:
+            raise ImportError(
+                inspect.getouterframes(inspect.currentframe())[1][3]
+                + ": shapefile version 2 or later required; try "
+                "pip install --upgrade pyshp"
+            )
+    return shapefile
 
 
 def write_gridlines_shapefile(filename, mg):
@@ -455,15 +480,10 @@ def shp2recarray(shpname):
     recarray : np.recarray
 
     """
-    try:
-        import shapefile as sf
-    except Exception:
-        raise Exception(
-            "io.to_shapefile(): error "
-            + "importing shapefile - try pip install pyshp"
-        )
     from ..utils.geometry import shape
 
+    sf = import_shapefile(check_version=False)
+
     sfobj = sf.Reader(shpname)
     dtype = [
         (str(f[0]), get_pyshp_field_dtypes(f[1])) for f in sfobj.fields[1:]
diff --git a/flopy/utils/geometry.py b/flopy/utils/geometry.py
index 713f6e668c..fc8633a77c 100644
--- a/flopy/utils/geometry.py
+++ b/flopy/utils/geometry.py
@@ -96,7 +96,7 @@ def pyshp_parts(self):
         # exterior ring must be clockwise (negative area)
         # interiors rings must be counter-clockwise (positive area)
 
-        shapefile = import_shapefile()
+        shapefile = import_shapefile(check_version=False)
 
         exterior = list(self.exterior)
         if shapefile.signed_area(exterior) > 0:
diff --git a/flopy/utils/gridgen.py b/flopy/utils/gridgen.py
index 956d176217..5cd3deb8e6 100644
--- a/flopy/utils/gridgen.py
+++ b/flopy/utils/gridgen.py
@@ -7,11 +7,8 @@
 from ..modflow.mfdisu import ModflowDisU
 from ..mf6.modflow import ModflowGwfdis
 from .util_array import Util2d  # read1d,
-from ..export.shapefile_utils import shp2recarray
+from ..export.shapefile_utils import import_shapefile, shp2recarray
 from ..mbase import which
-from ..export.shapefile_utils import import_shapefile
-
-shapefile = import_shapefile()
 
 
 # todo
@@ -54,6 +51,7 @@ def features_to_shapefile(features, featuretype, filename):
     None
 
     """
+    shapefile = import_shapefile(check_version=True)
 
     if featuretype.lower() not in ["point", "line", "polygon"]:
         raise Exception("Unrecognized feature type: {}".format(featuretype))
@@ -1899,6 +1897,7 @@ def _mkvertdict(self):
         None
 
         """
+        shapefile = import_shapefile(check_version=False)
         # ensure there are active leaf cells from gridgen
         fname = os.path.join(self.model_ws, "qtg.nod")
         if not os.path.isfile(fname):