Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions tests/recipes/test_write.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,26 @@ def test_specify_not_columns_regex():
df = wrangles.recipe.run(recipe, dataframe=data)
assert df.columns == ['3col']


def test_specify_not_columns_missing_ignored():
"""
Test writing and excluding a subset of columns while ignoring names that don't exist
"""
data = pd.DataFrame({
'col1': ['val1'],
'col2': ['val2']
})
recipe = """
write:
- dataframe:
not_columns:
- missing_col
- col2
"""
df = wrangles.recipe.run(recipe, dataframe=data)
assert df.columns.tolist() == ['col1']


def test_specify_where():
"""
Test writing and applying a WHERE sql criteria
Expand Down
35 changes: 30 additions & 5 deletions wrangles/recipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -1104,12 +1104,37 @@ def _filter_dataframe(
columns = _wildcard_expansion(df.columns.tolist(), columns)
df = df[columns]

# Remove any columns specified by the user
# Remove any columns specified by the user. Missing names are ignored so
# not_columns behaves like a no-op for absent fields instead of raising.
if not_columns:
not_columns = _wildcard_expansion(df.columns.tolist(), not_columns)
# List comprehension is used below to preserve order of columns
remaining_columns = [column for column in list(df.columns) if column not in not_columns]
df = df[remaining_columns]
if isinstance(not_columns, (str, int)):
not_columns = [not_columns]
else:
not_columns = list(not_columns)

valid_not_columns = []
for column in not_columns:
if column is None:
continue

# Negative selectors are valid even when they don't match any current
# columns, because they are meant to remove matching columns if present.
if isinstance(column, str) and column.startswith('-'):
valid_not_columns.append(column)
continue

try:
_wildcard_expansion(df.columns.tolist(), [column])
except KeyError:
continue

valid_not_columns.append(column)

if valid_not_columns:
not_columns = _wildcard_expansion(df.columns.tolist(), valid_not_columns)
# List comprehension is used below to preserve order of columns
remaining_columns = [column for column in list(df.columns) if column not in not_columns]
df = df[remaining_columns]

return df

Expand Down