Skip to content

Commit

Permalink
Use get to avoid nesting ifs
Browse files Browse the repository at this point in the history
  • Loading branch information
lajohn4747 committed Nov 22, 2023
1 parent 3dd9750 commit 0b4b42c
Showing 1 changed file with 11 additions and 16 deletions.
27 changes: 11 additions & 16 deletions sdmetrics/single_table/detection/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,22 +52,17 @@ def _drop_non_compute_columns(real_data, synthetic_data, metadata):
if metadata is not None:
drop_columns = []
drop_columns.extend(get_alternate_keys(metadata))
if 'columns' in metadata:
for column in metadata['columns']:
if ('primary_key' in metadata and
(column == metadata['primary_key'] or
column in metadata['primary_key'])):
drop_columns.append(column)

for field in metadata['columns'][column]:
if field == 'sdtype':
sdtype = metadata['columns'][column][field]
if sdtype not in ['numerical', 'datetime', 'categorical']:
drop_columns.append(column)

if field == 'pii':
if metadata['columns'][column][field]:
drop_columns.append(column)
for column in metadata.get('columns', []):
if ('primary_key' in metadata and
(column == metadata['primary_key'] or
column in metadata['primary_key'])):
drop_columns.append(column)

column_info = metadata['columns'].get(column, {})
sdtype = column_info.get('sdtype')
pii = column_info.get('pii')
if sdtype not in ['numerical', 'datetime', 'categorical'] or pii:
drop_columns.append(column)

if drop_columns:
transformed_real_data = real_data.drop(drop_columns, axis=1)
Expand Down

0 comments on commit 0b4b42c

Please sign in to comment.