add exists_ok param to add_column

parsons/etl/etl.py

@@ -1,6 +1,7 @@
-import petl
 import logging
 
+import petl
+
 logger = logging.getLogger(__name__)
 
 
@@ -9,7 +10,7 @@ def __init__(self):
 
         pass
 
-    def add_column(self, column, value=None, index=None):
+    def add_column(self, column, value=None, index=None, if_exists="fail"):
         """
         Add a column to your table
 
@@ -20,12 +21,19 @@ def add_column(self, column, value=None, index=None):
                 A fixed or calculated value
             index: int
                 The position of the new column in the table
+            if_exists: str (options: 'fail', 'replace')
+                If set `replace`, this function will call `fill_column`
+                if the column already exists, rather than raising a `ValueError`
         `Returns:`
             `Parsons Table` and also updates self
         """
 
         if column in self.columns:
-            raise ValueError(f"Column {column} already exists")
+            if if_exists == "replace":
+                self.fill_column(column, value)
+                return self
+            else:
+                raise ValueError(f"Column {column} already exists")
 
         self.table = self.table.addfield(column, value, index)
 

test/test_etl.py

@@ -1,13 +1,12 @@
-import unittest
-import petl
 import os
 import shutil
+import unittest
+from test.utils import assert_matching_tables
 
+import petl
 from parsons import Table
-from test.utils import assert_matching_tables
 from parsons.utilities import zip_archive
 
-
 # Notes :
 # - The `Table.to_postgres()` test is housed in the Postgres tests
 # - The `Table.from_postgres()` test is housed in the Postgres test
@@ -306,6 +305,10 @@ def test_column_add_dupe(self):
         # Test that we can't add an existing column name
         self.assertRaises(ValueError, self.tbl.add_column, "first")
 
+    def test_add_column_if_exists(self):
+        self.tbl.add_column("first", if_exists="replace")
+        self.assertEqual(self.tbl.columns, ["first", "last"])
+
     def test_remove_column(self):
         # Test that column is removed correctly
         self.tbl.remove_column("first")