Skip to content

Commit

Permalink
Remove external region columns
Browse files Browse the repository at this point in the history
  • Loading branch information
olejandro committed Feb 15, 2024
1 parent 98f7e45 commit 79b5a9d
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
1 change: 1 addition & 0 deletions xl2times/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ def convert_xl_to_times(
transforms.remove_comment_rows(config, t, model) for t in tables
],
transforms.process_regions,
transforms.remove_exreg_cols,
transforms.generate_dummy_processes,
transforms.process_time_slices,
transforms.process_transform_insert_variants,
Expand Down
43 changes: 43 additions & 0 deletions xl2times/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,49 @@ def remove_comment_cols(table: datatypes.EmbeddedXlTable) -> datatypes.EmbeddedX
return replace(table, dataframe=df)


def remove_exreg_cols(
config: datatypes.Config,
tables: List[datatypes.EmbeddedXlTable],
model: datatypes.TimesModel,
) -> List[datatypes.EmbeddedXlTable]:
"""
Remove external region columns from the all tables except tradelinks.
"""

external_regions = model.external_regions()

def remove_table_exreg_cols(
table: datatypes.EmbeddedXlTable,
) -> datatypes.EmbeddedXlTable:
"""
Return a modified copy of 'table' where columns that are external regions
have been removed.
"""

exreg_cols = [
colname
for colname in table.dataframe.columns
if colname.upper() in external_regions
]

if exreg_cols:
df = table.dataframe.drop(exreg_cols, axis=1)
return replace(table, dataframe=df)

else:
return table

# Do not do anything if external_reagions is empty
if not external_regions:
return tables
# Otherwise remove external region column from the relevant tables
else:
return [
remove_table_exreg_cols(t) if t.tag != datatypes.Tag.tradelinks else t
for t in tables
]


def remove_tables_with_formulas(
config: datatypes.Config,
tables: List[datatypes.EmbeddedXlTable],
Expand Down

0 comments on commit 79b5a9d

Please sign in to comment.