|
| 1 | +import shutil |
| 2 | +from pathlib import Path |
| 3 | + |
| 4 | +import pandas as pd |
| 5 | +from lir.aggregation import Aggregation, ContextAwareDict, config_parser, partial, pop_field |
| 6 | +from lir.util import check_type |
| 7 | + |
| 8 | + |
| 9 | +class CopyCSV(Aggregation): |
| 10 | + """Aggregation that copies a CSV file from a source location to a target location, optionally selecting columns. |
| 11 | +
|
| 12 | + Attributes |
| 13 | + ---------- |
| 14 | + source_file (str): The path to the source CSV file that should be copied. |
| 15 | + target_dir (str): The directory where the new CSV file will be saved. Given by the config_parser, meaning |
| 16 | + it is not set by the user in the configuration. |
| 17 | + columns (list[str]): A list of column names to copy from the source CSV. If empty, all columns will be copied. |
| 18 | + new_file_name (str): The name of the new CSV file. If empty, the original file name will be used. |
| 19 | + """ |
| 20 | + |
| 21 | + def __init__(self, source_file: str, target_dir: str, columns: list[str], new_file_name: str | None): |
| 22 | + self.source_file = Path(source_file) |
| 23 | + self.target_dir = Path(target_dir) |
| 24 | + self.columns = columns |
| 25 | + if new_file_name is None: |
| 26 | + self.new_file_name = self.target_dir / self.source_file.name |
| 27 | + else: |
| 28 | + self.new_file_name = self.target_dir / new_file_name |
| 29 | + |
| 30 | + def report(self, data) -> None: |
| 31 | + """Do nothing. Required by parent class.""" |
| 32 | + pass |
| 33 | + |
| 34 | + def close(self): |
| 35 | + """Close the aggregation and perform any necessary cleanup. |
| 36 | +
|
| 37 | + This method has the logic of this class. It copies the CSV file from the source to the target location, |
| 38 | + optionally selecting specific columns if they are specified. |
| 39 | + """ |
| 40 | + if self.columns: |
| 41 | + df = pd.read_csv(self.source_file) |
| 42 | + df[self.columns].to_csv(self.new_file_name, index=False) |
| 43 | + else: |
| 44 | + shutil.copy(self.source_file, self.new_file_name) |
| 45 | + |
| 46 | + |
| 47 | +@config_parser() |
| 48 | +def copy_csv(config: ContextAwareDict, output_dir: str) -> CopyCSV: |
| 49 | + """Parse the configuration for the CopyCSV aggregation and return an instance of it.""" |
| 50 | + source_file = pop_field(config, "file") |
| 51 | + columns = pop_field(config, "columns", default=[], validate=partial(check_type, list)) |
| 52 | + new_file_name = pop_field(config, "new_file_name", required=False) |
| 53 | + return CopyCSV(source_file, output_dir, columns=columns, new_file_name=new_file_name) |
0 commit comments