Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit d094d7d

Browse files
karimbahgatkarimbahgat
committed
Switch to fully streamed writer
With updated readme
1 parent f115b58 commit d094d7d

23 files changed

+258
-178
lines changed

‎README.md

Lines changed: 95 additions & 83 deletions
Large diffs are not rendered by default.

‎onlydbf.dbf

33 Bytes
Binary file not shown.

‎shapefile.dbf

33 Bytes
Binary file not shown.

‎shapefile.py

Lines changed: 163 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -874,18 +874,28 @@ def iterShapeRecords(self):
874874

875875
class Writer(object):
876876
"""Provides write support for ESRI Shapefiles."""
877-
def __init__(self, shapeType=None, autoBalance=False, bufsize=None, **kwargs):
877+
def __init__(self, target=None, shapeType=None, autoBalance=False, **kwargs):
878878
self.autoBalance = autoBalance
879879
self.fields = []
880880
self.shapeType = shapeType
881-
self.shp = None
882-
self.shx = None
883-
self.dbf = None
884-
# Create temporary files for immediate writing, minus the header
885-
self.bufsize = bufsize or 1056*1000*100 # default is 100 mb
886-
self._shp = tempfile.TemporaryFile()
887-
self._shx = tempfile.TemporaryFile()
888-
self._dbf = tempfile.TemporaryFile()
881+
self.shp = self.shx = self.dbf = None
882+
if target:
883+
self.shp = self.__getFileObj(os.path.splitext(target)[0] + '.shp')
884+
self.shx = self.__getFileObj(os.path.splitext(target)[0] + '.shx')
885+
self.dbf = self.__getFileObj(os.path.splitext(target)[0] + '.dbf')
886+
elif kwargs.get('shp') or kwargs.get('shx') or kwargs.get('dbf'):
887+
shp,shx,dbf = kwargs.get('shp'), kwargs.get('shx'), kwargs.get('dbf')
888+
if shp:
889+
self.shp = self.__getFileObj(shp)
890+
elif shx:
891+
self.shx = self.__getFileObj(shx)
892+
elif dbf:
893+
self.dbf = self.__getFileObj(dbf)
894+
else:
895+
raise Exception('Either the target filepath, or any of shp, shx, or dbf must be set to create a shapefile.')
896+
# Initiate with empty headers, to be finalized upon closing
897+
if self.shp: self.shp.write(b'9'*100)
898+
if self.shx: self.shx.write(b'9'*100)
889899
# Geometry record offsets and lengths for writing shx file.
890900
self.recNum = 0
891901
self.shpNum = 0
@@ -904,6 +914,56 @@ def __len__(self):
904914
of the two."""
905915
return max(self.recNum, self.shpNum)
906916

917+
def __enter__(self):
918+
"""
919+
Enter phase of context manager.
920+
"""
921+
return self
922+
923+
def __exit__(self, exc_type, exc_val, exc_tb):
924+
"""
925+
Exit phase of context manager, finish writing and close the files.
926+
"""
927+
self.close()
928+
929+
def __del__(self):
930+
self.close()
931+
932+
def close(self):
933+
"""
934+
Write final shp, shx, and dbf headers, close opened files.
935+
"""
936+
# Check if any of the files have already been closed
937+
shp_open = self.shp and not (hasattr(self.shp, 'closed') and self.shp.closed)
938+
shx_open = self.shx and not (hasattr(self.shx, 'closed') and self.shx.closed)
939+
dbf_open = self.dbf and not (hasattr(self.dbf, 'closed') and self.dbf.closed)
940+
941+
# Balance if already not balanced
942+
if self.shp and shp_open and self.dbf and dbf_open:
943+
if self.autoBalance:
944+
self.balance()
945+
if self.recNum != self.shpNum:
946+
raise ShapefileException("When saving both the dbf and shp file, "
947+
"the number of records (%s) must correspond "
948+
"with the number of shapes (%s)" % (self.recNum, self.shpNum))
949+
# Fill in the blank headers
950+
if self.shp and shp_open:
951+
self.__shapefileHeader(self.shp, headerType='shp')
952+
if self.shx and shx_open:
953+
self.__shapefileHeader(self.shx, headerType='shx')
954+
955+
# Update the dbf header with final length etc
956+
if self.dbf and dbf_open:
957+
self.__dbfHeader()
958+
959+
# Close files
960+
for attribute in (self.shp, self.shx, self.dbf):
961+
if hasattr(attribute, 'close'):
962+
try:
963+
attribute.close()
964+
except IOError:
965+
pass
966+
907967
def __getFileObj(self, f):
908968
"""Safety handler to verify file-like objects"""
909969
if not f:
@@ -914,17 +974,19 @@ def __getFileObj(self, f):
914974
pth = os.path.split(f)[0]
915975
if pth and not os.path.exists(pth):
916976
os.makedirs(pth)
917-
return open(f, "wb")
977+
return open(f, "wb+")
918978

919979
def __shpFileLength(self):
920980
"""Calculates the file length of the shp file."""
921-
# Start with header length
922-
size = 100
981+
# Remember starting position
982+
start = self.shp.tell()
923983
# Calculate size of all shapes
924-
self._shp.seek(0,2)
925-
size += self._shp.tell()
984+
self.shp.seek(0,2)
985+
size = self.shp.tell()
926986
# Calculate size as 16-bit words
927987
size //= 2
988+
# Return to start
989+
self.shp.seek(start)
928990
return size
929991

930992
def __bbox(self, s):
@@ -1116,8 +1178,8 @@ def shape(self, s):
11161178
self.__shxRecord(offset, length)
11171179

11181180
def __shpRecord(self, s):
1119-
f = self.__getFileObj(self._shp)
1120-
offset = 100 + f.tell()
1181+
f = self.__getFileObj(self.shp)
1182+
offset = f.tell()
11211183
# Record number, Content length place holder
11221184
self.shpNum += 1
11231185
f.write(pack(">2i", self.shpNum, 0))
@@ -1260,7 +1322,7 @@ def __shpRecord(self, s):
12601322

12611323
def __shxRecord(self, offset, length):
12621324
"""Writes the shx records."""
1263-
f = self.__getFileObj(self._shx)
1325+
f = self.__getFileObj(self.shx)
12641326
f.write(pack(">i", offset // 2))
12651327
f.write(pack(">i", length))
12661328

@@ -1297,7 +1359,13 @@ def record(self, *recordList, **recordDict):
12971359

12981360
def __dbfRecord(self, record):
12991361
"""Writes the dbf records."""
1300-
f = self.__getFileObj(self._dbf)
1362+
f = self.__getFileObj(self.dbf)
1363+
if self.recNum == 0:
1364+
# first records, so all fields should be set
1365+
# allowing us to write the dbf header
1366+
# cannot change the fields after this point
1367+
self.__dbfHeader()
1368+
# begin
13011369
self.recNum += 1
13021370
if not self.fields[0][0].startswith("Deletion"):
13031371
f.write(b' ') # deletion flag
@@ -1529,83 +1597,83 @@ def field(self, name, fieldType="C", size="50", decimal=0):
15291597
"Shapefile Writer reached maximum number of fields: 2046.")
15301598
self.fields.append((name, fieldType, size, decimal))
15311599

1532-
def saveShp(self, target):
1533-
"""Save an shp file."""
1534-
if not hasattr(target, "write"):
1535-
target = os.path.splitext(target)[0] + '.shp'
1536-
self.shp = self.__getFileObj(target)
1537-
self.__shapefileHeader(self.shp, headerType='shp')
1538-
self.shp.seek(100)
1539-
self._shp.seek(0)
1540-
chunk = True
1541-
while chunk:
1542-
chunk = self._shp.read(self.bufsize)
1543-
self.shp.write(chunk)
1544-
1545-
def saveShx(self, target):
1546-
"""Save an shx file."""
1547-
if not hasattr(target, "write"):
1548-
target = os.path.splitext(target)[0] + '.shx'
1549-
self.shx = self.__getFileObj(target)
1550-
self.__shapefileHeader(self.shx, headerType='shx')
1551-
self.shx.seek(100)
1552-
self._shx.seek(0)
1553-
chunk = True
1554-
while chunk:
1555-
chunk = self._shx.read(self.bufsize)
1556-
self.shx.write(chunk)
1557-
1558-
def saveDbf(self, target):
1559-
"""Save a dbf file."""
1560-
if not hasattr(target, "write"):
1561-
target = os.path.splitext(target)[0] + '.dbf'
1562-
self.dbf = self.__getFileObj(target)
1563-
self.__dbfHeader() # writes to .dbf
1564-
self._dbf.seek(0)
1565-
chunk = True
1566-
while chunk:
1567-
chunk = self._dbf.read(self.bufsize)
1568-
self.dbf.write(chunk)
1569-
1570-
def save(self, target=None, shp=None, shx=None, dbf=None):
1571-
"""Save the shapefile data to three files or
1572-
three file-like objects. SHP and DBF files can also
1573-
be written exclusively using saveShp, saveShx, and saveDbf respectively.
1574-
If target is specified but not shp, shx, or dbf then the target path and
1575-
file name are used. If no options or specified, a unique base file name
1576-
is generated to save the files and the base file name is returned as a
1577-
string.
1578-
"""
1579-
# Balance if already not balanced
1580-
if shp and dbf:
1581-
if self.autoBalance:
1582-
self.balance()
1583-
if self.recNum != self.shpNum:
1584-
raise ShapefileException("When saving both the dbf and shp file, "
1585-
"the number of records (%s) must correspond "
1586-
"with the number of shapes (%s)" % (self.recNum, self.shpNum))
1587-
# Save
1588-
if shp:
1589-
self.saveShp(shp)
1590-
if shx:
1591-
self.saveShx(shx)
1592-
if dbf:
1593-
self.saveDbf(dbf)
1594-
# Create a unique file name if one is not defined
1595-
if not shp and not shx and not dbf:
1596-
generated = False
1597-
if not target:
1598-
temp = tempfile.NamedTemporaryFile(prefix="shapefile_",dir=os.getcwd())
1599-
target = temp.name
1600-
generated = True
1601-
self.saveShp(target)
1602-
self.shp.close()
1603-
self.saveShx(target)
1604-
self.shx.close()
1605-
self.saveDbf(target)
1606-
self.dbf.close()
1607-
if generated:
1608-
return target
1600+
## def saveShp(self, target):
1601+
## """Save an shp file."""
1602+
## if not hasattr(target, "write"):
1603+
## target = os.path.splitext(target)[0] + '.shp'
1604+
## self.shp = self.__getFileObj(target)
1605+
## self.__shapefileHeader(self.shp, headerType='shp')
1606+
## self.shp.seek(100)
1607+
## self._shp.seek(0)
1608+
## chunk = True
1609+
## while chunk:
1610+
## chunk = self._shp.read(self.bufsize)
1611+
## self.shp.write(chunk)
1612+
##
1613+
## def saveShx(self, target):
1614+
## """Save an shx file."""
1615+
## if not hasattr(target, "write"):
1616+
## target = os.path.splitext(target)[0] + '.shx'
1617+
## self.shx = self.__getFileObj(target)
1618+
## self.__shapefileHeader(self.shx, headerType='shx')
1619+
## self.shx.seek(100)
1620+
## self._shx.seek(0)
1621+
## chunk = True
1622+
## while chunk:
1623+
## chunk = self._shx.read(self.bufsize)
1624+
## self.shx.write(chunk)
1625+
##
1626+
## def saveDbf(self, target):
1627+
## """Save a dbf file."""
1628+
## if not hasattr(target, "write"):
1629+
## target = os.path.splitext(target)[0] + '.dbf'
1630+
## self.dbf = self.__getFileObj(target)
1631+
## self.__dbfHeader() # writes to .dbf
1632+
## self._dbf.seek(0)
1633+
## chunk = True
1634+
## while chunk:
1635+
## chunk = self._dbf.read(self.bufsize)
1636+
## self.dbf.write(chunk)
1637+
1638+
## def save(self, target=None, shp=None, shx=None, dbf=None):
1639+
## """Save the shapefile data to three files or
1640+
## three file-like objects. SHP and DBF files can also
1641+
## be written exclusively using saveShp, saveShx, and saveDbf respectively.
1642+
## If target is specified but not shp, shx, or dbf then the target path and
1643+
## file name are used. If no options or specified, a unique base file name
1644+
## is generated to save the files and the base file name is returned as a
1645+
## string.
1646+
## """
1647+
## # Balance if already not balanced
1648+
## if shp and dbf:
1649+
## if self.autoBalance:
1650+
## self.balance()
1651+
## if self.recNum != self.shpNum:
1652+
## raise ShapefileException("When saving both the dbf and shp file, "
1653+
## "the number of records (%s) must correspond "
1654+
## "with the number of shapes (%s)" % (self.recNum, self.shpNum))
1655+
## # Save
1656+
## if shp:
1657+
## self.saveShp(shp)
1658+
## if shx:
1659+
## self.saveShx(shx)
1660+
## if dbf:
1661+
## self.saveDbf(dbf)
1662+
## # Create a unique file name if one is not defined
1663+
## if not shp and not shx and not dbf:
1664+
## generated = False
1665+
## if not target:
1666+
## temp = tempfile.NamedTemporaryFile(prefix="shapefile_",dir=os.getcwd())
1667+
## target = temp.name
1668+
## generated = True
1669+
## self.saveShp(target)
1670+
## self.shp.close()
1671+
## self.saveShx(target)
1672+
## self.shx.close()
1673+
## self.saveDbf(target)
1674+
## self.dbf.close()
1675+
## if generated:
1676+
## return target
16091677

16101678
# Begin Testing
16111679
def test(**kwargs):

‎shapefile.pyc

-615 Bytes
Binary file not shown.

‎shapefile.shp

100 Bytes
Binary file not shown.

‎shapefile.shx

100 Bytes
Binary file not shown.

‎shapefiles/test/balancing.dbf

804 Bytes
Binary file not shown.

‎shapefiles/test/balancing.shp

264 Bytes
Binary file not shown.

‎shapefiles/test/balancing.shx

156 Bytes
Binary file not shown.

‎shapefiles/test/contextwriter.dbf

33 Bytes
Binary file not shown.

‎shapefiles/test/contextwriter.shp

100 Bytes
Binary file not shown.

‎shapefiles/test/contextwriter.shx

100 Bytes
Binary file not shown.

‎shapefiles/test/dtype.dbf

182 Bytes
Binary file not shown.

‎shapefiles/test/dtype.shp

-48 Bytes
Binary file not shown.

‎shapefiles/test/dtype.shx

-32 Bytes
Binary file not shown.

‎shapefiles/test/line.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/linem.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/linez.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/multipatch.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/multipoint.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/point.dbf

0 Bytes
Binary file not shown.

‎shapefiles/test/polygon.dbf

0 Bytes
Binary file not shown.

0 commit comments

Comments
 (0)