Store count of shp shapes, fixing misc bugs if dbf is absent, add tests for loading shp/shx/dbf separately

Author: karimbahgat

shapefile.py

@@ -796,6 +796,7 @@ def __init__(self, *args, **kwargs):
         self._offsets = []
         self.shpLength = None
         self.numRecords = None
+        self.numShapes = None
         self.fields = []
         self.__dbfHdrLength = 0
         self.__fieldposition_lookup = {}
@@ -862,7 +863,38 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def __len__(self):
         """Returns the number of shapes/records in the shapefile."""
-        return self.numRecords
+        if self.dbf:
+            # Preferably use dbf record count
+            if self.numRecords is None:
+                self.__dbfHeader()
+                
+            return self.numRecords
+        
+        elif self.shp:
+            # Otherwise use shape count
+            if self.shx:
+                # Use index file to get total count
+                if self.numShapes is None:
+                    # File length (16-bit word * 2 = bytes) - header length
+                    self.shx.seek(24)
+                    shxRecordLength = (unpack(">i", self.shx.read(4))[0] * 2) - 100
+                    self.numShapes = shxRecordLength // 8
+                    
+                return self.numShapes
+            
+            else:
+                # Index file not available, iterate all shapes to get total count
+                if self.numShapes is None:
+                    i = 0
+                    for i,shape in enumerate(self.iterShapes()):
+                        i += 1
+                    self.numShapes = i
+                    
+                return self.numShapes
+            
+        else:
+            # No file loaded yet, treat as 'empty' shapefile
+            return 0 
 
     def __iter__(self):
         """Iterates through the shapes/records in the shapefile."""
@@ -1072,15 +1104,16 @@ def __shapeIndex(self, i=None):
         if not shx:
             return None
         if not self._offsets:
-            # File length (16-bit word * 2 = bytes) - header length
-            shx.seek(24)
-            shxRecordLength = (unpack(">i", shx.read(4))[0] * 2) - 100
-            numRecords = shxRecordLength // 8
+            if self.numShapes is None:
+                # File length (16-bit word * 2 = bytes) - header length
+                shx.seek(24)
+                shxRecordLength = (unpack(">i", shx.read(4))[0] * 2) - 100
+                self.numShapes = shxRecordLength // 8
             # Jump to the first record.
             shx.seek(100)
             shxRecords = _Array('i')
             # Each offset consists of two nrs, only the first one matters
-            shxRecords.fromfile(shx, 2 * numRecords)
+            shxRecords.fromfile(shx, 2 * self.numShapes)
             if sys.byteorder != 'big':
                  shxRecords.byteswap()
             self._offsets = [2 * el for el in shxRecords[::2]]

test_shapefile.py

@@ -276,6 +276,93 @@ def test_reader_fields():
         assert isinstance(field[3], int)    # decimal length
 
 
+def test_reader_shapefile_extension_ignored():
+    """
+    Assert that the filename's extension is
+    ignored when reading a shapefile.
+    """
+    base = "shapefiles/blockgroups"
+    ext = ".abc"
+    filename = base + ext
+    with shapefile.Reader(filename) as sf:
+        assert len(sf) == 663
+
+    # assert test.abc does not exist
+    assert not os.path.exists(filename)
+
+
+def test_reader_dbf_only():
+    """
+    Assert that specifying just the
+    dbf argument to the shapefile reader
+    reads just the dbf file.
+    """
+    with shapefile.Reader(dbf="shapefiles/blockgroups.dbf") as sf:
+        assert len(sf) == 663
+        record = sf.record(3)
+        assert record[1:3] == ['060750601001', 4715]
+
+
+def test_reader_shp_shx_only():
+    """
+    Assert that specifying just the
+    shp and shx argument to the shapefile reader
+    reads just the shp and shx file.
+    """
+    with shapefile.Reader(shp="shapefiles/blockgroups.shp", shx="shapefiles/blockgroups.shx") as sf:
+        assert len(sf) == 663
+        shape = sf.shape(3)
+        assert len(shape.points) is 173
+
+
+def test_reader_shx_optional():
+    """
+    Assert that specifying just the
+    shp argument to the shapefile reader
+    reads just the shp file (shx optional).
+    """
+    with shapefile.Reader(shp="shapefiles/blockgroups.shp") as sf:
+        assert len(sf) == 663
+        shape = sf.shape(3)
+        assert len(shape.points) is 173
+
+
+def test_reader_filelike_dbf_only():
+    """
+    Assert that specifying just the
+    dbf argument to the shapefile reader
+    reads just the dbf file.
+    """
+    with shapefile.Reader(dbf=open("shapefiles/blockgroups.dbf", "rb")) as sf:
+        assert len(sf) == 663
+        record = sf.record(3)
+        assert record[1:3] == ['060750601001', 4715]
+
+
+def test_reader_filelike_shp_shx_only():
+    """
+    Assert that specifying just the
+    shp and shx argument to the shapefile reader
+    reads just the shp and shx file.
+    """
+    with shapefile.Reader(shp=open("shapefiles/blockgroups.shp", "rb"), shx=open("shapefiles/blockgroups.shx", "rb")) as sf:
+        assert len(sf) == 663
+        shape = sf.shape(3)
+        assert len(shape.points) is 173
+
+
+def test_reader_filelike_shx_optional():
+    """
+    Assert that specifying just the
+    shp argument to the shapefile reader
+    reads just the shp file (shx optional).
+    """
+    with shapefile.Reader(shp=open("shapefiles/blockgroups.shp", "rb")) as sf:
+        assert len(sf) == 663
+        shape = sf.shape(3)
+        assert len(shape.points) is 173
+
+
 def test_records_match_shapes():
     """
     Assert that the number of records matches
@@ -498,7 +585,7 @@ def test_write_geojson(tmpdir):
 @pytest.mark.parametrize("shape_type", shape_types)
 def test_write_empty_shapefile(tmpdir, shape_type):
     """
-    Assert that can write an empty shapefile
+    Assert that can write an empty shapefile, for all different shape types. 
     """
     filename = tmpdir.join("test").strpath
     with shapefile.Writer(filename, shapeType=shape_type) as w: