Merge pull request #61 from ckirby/issue-50

Add with_header boolean kwarg to .to_csv()

docs/index.rst

@@ -444,7 +444,7 @@ Export options
 
 The ``to_csv`` manager method only requires one argument, the path to where the CSV should be exported. It also allows users to optionally limit or expand the fields written out by providing them as additional parameters.
 
-.. method:: to_csv(csv_path [, *fields, delimiter=','])
+.. method:: to_csv(csv_path [, *fields, delimiter=',', with_header=True])
 
 
 =================  =========================================================
@@ -460,6 +460,8 @@ Argument           Description
 
 ``delimiter``      String that will be used as a delimiter for the CSV
                    file.
+
+``with_header``    Boolean determines if the header should be exported
 =================  =========================================================
 
 Reducing the exported fields

postgres_copy/__init__.py

@@ -25,6 +25,8 @@ def to_csv(self, csv_path, *fields, **kwargs):
         query = self.query.clone(CopyToQuery)
         query.copy_to_fields = fields
         query.copy_to_delimiter = kwargs.get('delimiter', ',')
+        with_header = kwargs.get('with_header', True)
+        query.copy_to_header = "HEADER" if with_header else ""
         compiler = query.get_compiler(self.db, connection=connection)
         compiler.execute_sql(csv_path)
 

postgres_copy/copy_to.py

@@ -49,8 +49,8 @@ def execute_sql(self, csv_path):
                 select_sql = self.as_sql()[0] % adapted_params
                 # then the COPY TO query
                 copy_to_sql = (
-                    "COPY (%s) TO STDOUT DELIMITER '%s' CSV HEADER"
-                ) % (select_sql, self.query.copy_to_delimiter)
+                    "COPY (%s) TO STDOUT DELIMITER '%s' CSV %s"
+                ) % (select_sql, self.query.copy_to_delimiter, self.query.copy_to_header)
                 # then execute
                 c.cursor.copy_expert(copy_to_sql, stdout)
 

tests/tests.py

@@ -63,6 +63,29 @@ def test_export(self):
             [i['name'] for i in reader]
         )
 
+    def test_export_header_setting(self):
+        self._load_objects(self.name_path)
+        MockObject.objects.to_csv(self.export_path)
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        self.assertTrue(
+            ['BEN', 'JOE', 'JANE'],
+            [i['name'] for i in reader]
+        )
+        MockObject.objects.to_csv(self.export_path, with_header=True)
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        self.assertTrue(
+            ['BEN', 'JOE', 'JANE'],
+            [i['name'] for i in reader]
+        )
+        MockObject.objects.to_csv(self.export_path, with_header=False)
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        with self.assertRaises(KeyError):
+            [i['name'] for i in reader]
+        self.assertTrue(
+            ['JOE', 'JANE'],
+            [i['BEN'] for i in reader]
+        )
+
     def test_export_delimiter(self):
         self._load_objects(self.name_path)
         MockObject.objects.to_csv(self.export_path, delimiter=';')