Seamlessly Read/Write datetime.date Objects

shapefile.py

@@ -15,6 +15,7 @@
 import time
 import array
 import tempfile
+from datetime import date
 
 #
 # Constants for shape types
@@ -473,8 +474,7 @@ def __record(self):
             # deleted record
             return None
         record = []
-        for (name, typ, size, deci), value in zip(self.fields,
-                                                                                                recordContents):
+        for (name, typ, size, deci), value in zip(self.fields, recordContents):
             if name == 'DeletionFlag':
                 continue
             elif not value.strip():
@@ -491,7 +491,7 @@ def __record(self):
             elif typ == b('D'):
                 try:
                     y, m, d = int(value[:4]), int(value[4:6]), int(value[6:8])
-                    value = [y, m, d]
+                    value = date(y, m, d)
                 except:
                     value = value.strip()
             elif typ == b('L'):
@@ -882,10 +882,12 @@ def __dbfRecords(self):
             for (fieldName, fieldType, size, dec), value in zip(self.fields, record):
                 fieldType = fieldType.upper()
                 size = int(size)
-                if fieldType.upper() == "N":
+                if fieldType == 'N':
                     value = str(value).rjust(size)
                 elif fieldType == 'L':
                     value = str(value)[0].upper()
+                elif fieldType == 'D' and isinstance(value, date):
+                    value = value.strftime('%Y%m%d').ljust(size)
                 else:
                     value = str(value)[:size].ljust(size)
                 assert len(value) == size
@@ -939,10 +941,19 @@ def poly(self, parts=[], shapeType=POLYGON, partTypes=[]):
             polyShape.partTypes = partTypes
         self._shapes.append(polyShape)
 
-    def field(self, name, fieldType="C", size="50", decimal=0):
+    def field(self, name, fieldType='C', size=50, decimal=0):
         """Adds a dbf field descriptor to the shapefile."""
+        fieldType = fieldType.upper()
+        size = int(size)
+        decimal = int(decimal)
+        if fieldType == 'D' and size < 8:
+            raise ShapefileException("Date fields require min size of 8 (%d is too small)" % size)
         self.fields.append((name, fieldType, size, decimal))
 
+    def fieldDate(self, name):
+        """Adds a Date type dbf descriptor to the shapefile"""
+        self.field(name, 'D', 8, 0)
+
     def record(self, *recordList, **recordDict):
         """Creates a dbf attribute record. You can submit either a sequence of
         field values or keyword arguments of field names and values. Before

tests/date_tests.py

@@ -0,0 +1,41 @@
+"""
+Unit tests for Pythonic `datetime.date` object handling.
+"""
+
+import unittest
+import datetime
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import BytesIO as StringIO
+
+import shapefile
+
+
+class TestDateHandling(unittest.TestCase):
+
+    def testDateReadWrite(self):
+        """Round-trip read and write Python `date` as value for 'D' field"""
+        today = datetime.date.today()
+
+        # Write a one-field, one-record shp to memory; use `date` obj as value
+        writer = shapefile.Writer()
+        writer.field('DATEFIELD', 'D', 8)
+        writer.null()
+        writer.record(today)
+        shp, shx, dbf = StringIO(), StringIO(), StringIO()
+        writer.save(shp=shp, shx=shx, dbf=dbf)
+
+        # Read our in-memory shp to verify that Reader gives us a `date` obj
+        reader = shapefile.Reader(shp=shp, shx=shx, dbf=dbf)
+        self.assertEqual(reader.fields[-1][1], 'D')
+        self.assertEqual(len(reader.records()), 1)
+        record = reader.record(0)
+        d = record[0]
+        self.assertTrue(isinstance(d, datetime.date),
+                        "Expected a `date` object back from Reader (we got a %s)" % type(d))
+        self.assertEqual(d, today)
+
+
+if __name__ == '__main__':
+    unittest.main()