diff --git a/parsons/etl/etl.py b/parsons/etl/etl.py index 75f9762fed..7797db0f2b 100644 --- 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 index 3a3013b522..d17b900bc3 100644 --- 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")