ops_utils.csv_util

  1"""Module for CSV file operations."""
  2import csv
  3import logging
  4from typing import Optional, Sequence
  5
  6
  7class Csv:
  8    """Class for CSV file operations."""
  9
 10    def __init__(self, file_path: str, delimiter: str = "\t"):
 11        r"""
 12        Initialize the Csv class.
 13
 14        **Args:**
 15        - file_path (str): The path to the tabular file.
 16        - delimiter (str, optional): The delimiter to use in the tabular file. Defaults to `\\t` (tab-delimited).
 17        """
 18        self.file_path = file_path
 19        """@private"""
 20        self.delimiter = delimiter
 21        """@private"""
 22
 23    def create_tsv_from_list_of_dicts(self, list_of_dicts: list[dict], header_list: Optional[list[str]] = None) -> str:
 24        """
 25        Create a TSV file from a list of dictionaries.
 26
 27        **Args:**
 28        - list_of_dicts (list[dict]): The list of dictionaries to write to the TSV file.
 29        - header_list (list[str], optional): The list of headers to use in the TSV file.
 30                If provided, output columns will be in same order as list. Defaults to None.
 31
 32        **Returns:**
 33        - str: The path to the created TSV file.
 34        """
 35        # Create one flat unique list by doing list comprehension where it loops
 36        # through twice to make it flat and transform to set and back to list
 37        # to make it unique
 38        if not header_list:
 39            header_list = sorted(
 40                list(
 41                    set(
 42                        [
 43                            header_list
 44                            for d in list_of_dicts
 45                            for header_list in d.keys()
 46                        ]
 47                    )
 48                )
 49            )
 50        logging.info(f'Creating {self.file_path}')
 51        with open(self.file_path, 'w', newline='') as f:
 52            writer = csv.DictWriter(
 53                f, fieldnames=header_list, delimiter='\t', quotechar="'", extrasaction='ignore')
 54            writer.writeheader()
 55            for d in list_of_dicts:
 56                writer.writerow(d)
 57        return self.file_path
 58
 59    def create_tsv_from_list_of_lists(self, list_of_lists: list[list]) -> str:
 60        """
 61        Create a TSV file from a list of lists.
 62
 63        **Args:**
 64        - list_of_lists (list[list]): The list of lists to write to the TSV file.
 65
 66        **Returns:**
 67        - str: The path to the created TSV file.
 68        """
 69        logging.info(f'Creating {self.file_path}')
 70        with open(self.file_path, 'w') as f:
 71            for list_of_data in list_of_lists:
 72                # Make sure all entries are strings
 73                str_only_list = [str(entry) for entry in list_of_data]
 74                f.write(self.delimiter.join(str_only_list) + '\n')
 75        return self.file_path
 76
 77    def create_list_of_dicts_from_tsv_with_no_headers(self, headers_list: list[str]) -> list[dict]:
 78        """
 79        Create a list of dictionaries from a TSV file with no headers.
 80
 81        **Args:**
 82        - headers_list (list[str]): The list of headers to use for the TSV file.
 83
 84        **Returns:**
 85        - list[dict]: The list of dictionaries created from the TSV file.
 86        """
 87        with open(self.file_path, 'r') as f:
 88            reader = csv.DictReader(
 89                f, delimiter=self.delimiter, fieldnames=headers_list)
 90            return [row for row in reader]
 91
 92    def get_header_order_from_tsv(self) -> Sequence[str]:
 93        """
 94        Get the header order from a TSV file.
 95
 96        **Returns:**
 97        - list[str]: The list of headers in the TSV file.
 98        """
 99        with open(self.file_path, 'r') as f:
100            reader = csv.DictReader(f, delimiter=self.delimiter, skipinitialspace=True)
101            return reader.fieldnames  # type: ignore[return-value]
102
103    def create_list_of_dicts_from_tsv(
104            self,
105            expected_headers: Optional[list[str]] = None,
106            allow_extra_headers: bool = False
107    ) -> list[dict]:
108        """
109        Create a list of dictionaries from a TSV file.
110
111        **Args:**
112        - expected_headers (list[str], optional): The list of expected headers. If provided,
113                will check that all headers are present in the TSV file. Defaults to None.
114        - allow_extra_headers (bool, optional): Whether to allow extra headers in the TSV file.
115                Only used if `expected_headers` is provided. Defaults to False.
116
117        **Returns:**
118        - list[dict]: The list of dictionaries created from the TSV file.
119
120        **Raises:**
121        - ValueError: If the expected headers are not found in the TSV file.
122        """
123        with open(self.file_path) as f:
124            dict_reader = csv.DictReader(
125                f, delimiter=self.delimiter, skipinitialspace=True)
126            if expected_headers:
127                match = True
128                tsv_headers = dict_reader.fieldnames
129                extra_headers = set(tsv_headers) - set(expected_headers)  # type: ignore[arg-type]
130                missing_headers = set(expected_headers) - set(tsv_headers)  # type: ignore[arg-type]
131                if extra_headers:
132                    extra_string = ','.join(extra_headers)
133                    logging.warning(
134                        f"Extra headers found in tsv: {extra_string}")
135                    if not allow_extra_headers:
136                        match = False
137                if missing_headers:
138                    missing_string = ','.join(missing_headers)
139                    logging.error(
140                        f"Missing expected headers: {missing_string}")
141                    match = False
142                if not match:
143                    raise ValueError(
144                        f"Expected headers not in {self.file_path}")
145            return [
146                {
147                    k: v
148                    for k, v in row.items()
149                }
150                for row in dict_reader
151            ]