accept issue 3436

1 files changed, 33 insertions, 0 deletions

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: