add exists_ok param to add_column (#839)

* add exists_ok param to add_column * feat(test): test exists_ok * chore(linting): Black parsons * switch exists_ok to if_exists * chore(test): Fix test to use new param * chore(linting): Lint tests too --------- Co-authored-by: Tal Levy <[email protected]>
move-coop · Jun 20, 2023 · e7f1e94 · e7f1e94
1 parent 71a9672
commit e7f1e94
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/parsons/etl/etl.py b/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)
 

diff --git a/test/test_etl.py b/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")