From a032bf41f6a23856f92690c0cd14520d31214e7e Mon Sep 17 00:00:00 2001
From: Skip Montanaro <skip@pobox.com>
Date: Fri, 8 Aug 2008 22:52:51 +0000
Subject: accept issue 3436

---
 Doc/library/csv.rst  | 12 ++++++++++++
 Lib/csv.py           | 22 ++++++++++++++++++----
 Lib/test/test_csv.py | 33 +++++++++++++++++++++++++++++++++
 Misc/NEWS            |  4 ++++
 4 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/Doc/library/csv.rst b/Doc/library/csv.rst
index eef3cfd..2458b3c 100644
--- a/Doc/library/csv.rst
+++ b/Doc/library/csv.rst
@@ -382,6 +382,18 @@ Reader objects have the following public attributes:
    .. versionadded:: 2.5
 
 
+DictReader objects have the following public attribute:
+
+
+.. attribute:: csvreader.fieldnames
+
+   If not passed as a parameter when creating the object, this attribute is
+   initialized upon first access or when the first record is read from the
+   file.
+
+   .. versionchanged:: 2.6
+
+
 Writer Objects
 --------------
 
diff --git a/Lib/csv.py b/Lib/csv.py
index 41026e0..7210dac 100644
--- a/Lib/csv.py
+++ b/Lib/csv.py
@@ -71,7 +71,7 @@ register_dialect("excel-tab", excel_tab)
 class DictReader:
     def __init__(self, f, fieldnames=None, restkey=None, restval=None,
                  dialect="excel", *args, **kwds):
-        self.fieldnames = fieldnames    # list of keys for the dict
+        self._fieldnames = fieldnames   # list of keys for the dict
         self.restkey = restkey          # key to catch long rows
         self.restval = restval          # default value for short rows
         self.reader = reader(f, dialect, *args, **kwds)
@@ -81,11 +81,25 @@ class DictReader:
     def __iter__(self):
         return self
 
+    @property
+    def fieldnames(self):
+        if self._fieldnames is None:
+            try:
+                self._fieldnames = self.reader.next()
+            except StopIteration:
+                pass
+        self.line_num = self.reader.line_num
+        return self._fieldnames
+
+    @fieldnames.setter
+    def fieldnames(self, value):
+        self._fieldnames = value
+
     def next(self):
+        if self.line_num == 0:
+            # Used only for its side effect.
+            self.fieldnames
         row = self.reader.next()
-        if self.fieldnames is None:
-            self.fieldnames = row
-            row = self.reader.next()
         self.line_num = self.reader.line_num
 
         # unlike the basic reader, we prefer not to return blanks,
diff --git a/Lib/test/test_csv.py b/Lib/test/test_csv.py
index 8a89f4c..b239f75 100644
--- a/Lib/test/test_csv.py
+++ b/Lib/test/test_csv.py
@@ -611,11 +611,43 @@ class TestDictFields(unittest.TestCase):
             fileobj.write("f1,f2,f3\r\n1,2,abc\r\n")
             fileobj.seek(0)
             reader = csv.DictReader(fileobj)
+            self.assertEqual(reader.fieldnames, ["f1", "f2", "f3"])
             self.assertEqual(reader.next(), {"f1": '1', "f2": '2', "f3": 'abc'})
         finally:
             fileobj.close()
             os.unlink(name)
 
+    # Two test cases to make sure existing ways of implicitly setting
+    # fieldnames continue to work.  Both arise from discussion in issue3436.
+    def test_read_dict_fieldnames_from_file(self):
+        fd, name = tempfile.mkstemp()
+        f = os.fdopen(fd, "w+b")
+        try:
+            f.write("f1,f2,f3\r\n1,2,abc\r\n")
+            f.seek(0)
+            reader = csv.DictReader(f, fieldnames=csv.reader(f).next())
+            self.assertEqual(reader.fieldnames, ["f1", "f2", "f3"])
+            self.assertEqual(reader.next(), {"f1": '1', "f2": '2', "f3": 'abc'})
+        finally:
+            f.close()
+            os.unlink(name)
+
+    def test_read_dict_fieldnames_chain(self):
+        import itertools
+        fd, name = tempfile.mkstemp()
+        f = os.fdopen(fd, "w+b")
+        try:
+            f.write("f1,f2,f3\r\n1,2,abc\r\n")
+            f.seek(0)
+            reader = csv.DictReader(f)
+            first = next(reader)
+            for row in itertools.chain([first], reader):
+                self.assertEqual(reader.fieldnames, ["f1", "f2", "f3"])
+                self.assertEqual(row, {"f1": '1', "f2": '2', "f3": 'abc'})
+        finally:
+            f.close()
+            os.unlink(name)
+
     def test_read_long(self):
         fd, name = tempfile.mkstemp()
         fileobj = os.fdopen(fd, "w+b")
@@ -651,6 +683,7 @@ class TestDictFields(unittest.TestCase):
             fileobj.write("f1,f2\r\n1,2,abc,4,5,6\r\n")
             fileobj.seek(0)
             reader = csv.DictReader(fileobj, restkey="_rest")
+            self.assertEqual(reader.fieldnames, ["f1", "f2"])
             self.assertEqual(reader.next(), {"f1": '1', "f2": '2',
                                              "_rest": ["abc", "4", "5", "6"]})
         finally:
diff --git a/Misc/NEWS b/Misc/NEWS
index 0ee5fd5..0989cfe 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -41,6 +41,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #3436: Make csv.DictReader's fieldnames attribute a property so that
+  upon first access it can be automatically initialized from the csv file if
+  it wasn't initialized during instantiation.
+
 - Issue #2338: Create imp.reload() to help with transitioning to Python 3.0 as
   the reload() built-in has been removed.
 
-- 
cgit v0.12