From 6de7f0bf8402a3a28df743e7df3403a228d3aaf1 Mon Sep 17 00:00:00 2001
From: Chapin Bryce <27cbryce@gmail.com>
Date: Mon, 27 May 2019 19:51:58 -0400
Subject: [PATCH] Added CSV writing code

---
 sections/section_01_essentials/csv_example.py | 64 +++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/sections/section_01_essentials/csv_example.py b/sections/section_01_essentials/csv_example.py
index 5d48ade..5db73ac 100644
--- a/sections/section_01_essentials/csv_example.py
+++ b/sections/section_01_essentials/csv_example.py
@@ -34,3 +34,67 @@ __docs__ = [
     'https://docs.python.org/3/library/os.html'
 ]
 
+
+def write_csv_dicts(outfile, data, headers=None):
+    """Writes a list of dictionaries to a CSV file.
+
+    Arguments:
+        outfile (str): Path to output file
+        data (list): List of dictionaries to write to file
+        headers (list): Header row to use. If empty, will use the
+            first dictionary in the `data` list.
+    """
+
+    if not headers:
+        # Use the first line of data
+        headers = [str(x) for x in data[0].keys()]
+
+    with open(outfile, 'w', newline="") as open_file:
+        # Write only provided headers, ignore others
+        csvfile = csv.DictWriter(open_file, headers,
+                                 extrasaction='ignore')
+        csvfile.writeheader()
+        csvfile.writerows(data)
+
+
+def write_csv_lists(outfile, data, headers=None):
+    """Writes a list of lists to a CSV file.
+
+    Arguments:
+        outfile (str): Path to output file
+        data (list): List of lists to write to file
+        headers (list): Header row to use. If empty, will use the
+            first list in the `data` list.
+    """
+
+    with open(outfile, 'w', newline="") as open_file:
+        # Write only provided headers, ignore others
+        csvfile = csv.writer(open_file)
+        for count, entry in enumerate(data):
+            if count == 0:
+                # If headers are defined, write them, otherwise
+                # continue as they will be written anyways
+                if headers:
+                    csvfile.writerow(headers)
+            csvfile.writerow(entry)
+
+sample_dict_data = [
+    {'id': '0', 'city': 'Boston', 'state': 'MA',
+     'country': 'USA'},
+    {'id': '1', 'city': 'New York', 'state': 'NY',
+     'country': 'USA'},
+    {'id': '2', 'city': 'Washington', 'state': 'DC',
+     'country': 'USA'},
+]
+
+write_csv_dicts('dict_test.csv', sample_dict_data)
+
+headers = ['id', 'city', 'state', 'country']
+sample_list_data = [
+    ['0', 'Boston', 'MA', 'USA'],
+    ['1', 'New York', 'NY', 'USA'],
+    ['2', 'Washington', 'DC', 'USA']
+]
+
+write_csv_lists('list_test.csv', sample_list_data,
+                headers=headers)