diff --git a/src/psPlotKit/data_manager/ps_data.py b/src/psPlotKit/data_manager/ps_data.py index c736485..a627a9e 100644 --- a/src/psPlotKit/data_manager/ps_data.py +++ b/src/psPlotKit/data_manager/ps_data.py @@ -169,13 +169,40 @@ def _define_custom_units(self, custom_units=None): custom_units = CustomUnits() self.custom_units = custom_units.get_units_dict() + def _raise_unit_error(self, error, operation, requested_units=None, other=None): + """Build and raise an enhanced unit error preserving the original type. + + The raised exception keeps the same Python type as *error* so existing + ``except ValueError`` / ``except LookupError`` blocks continue to work, + but the message now includes the data_key and data_type so users can + identify which data triggered the incompatible unit problem. + """ + msg = "Unit error for data_key={!r}, data_type={!r}. Operation: {}. ".format( + self.data_key, self.data_type, operation + ) + if requested_units is not None: + msg += "Requested units: {!r}. ".format(requested_units) + if other is not None: + if isinstance(other, PsData): + msg += "Other operand: data_key={!r}, data_type={!r}. ".format( + other.data_key, other.data_type + ) + else: + msg += "Other operand: {!r}. ".format(other) + msg += "Original error: {}: {}".format(type(error).__name__, error) + _logger.error(msg) + raise type(error)(msg) from error + def _assign_units(self, manual_conversion=1): if self.data_is_numbers: self.data = self.data * manual_conversion qsunits = self._get_qs_unit() - self.data_with_units = qs.Quantity(self.data.copy(), qsunits) - self.raw_data_with_units = qs.Quantity(self.raw_data.copy(), qsunits) - self.data = self.data_with_units.magnitude + try: + self.data_with_units = qs.Quantity(self.data.copy(), qsunits) + self.raw_data_with_units = qs.Quantity(self.raw_data.copy(), qsunits) + self.data = self.data_with_units.magnitude + except (ValueError, LookupError) as e: + self._raise_unit_error(e, "assign units", requested_units=self.sunits) self.set_label() @property @@ -280,15 +307,21 @@ def set_data(self, data): self.raw_data = data def to_units(self, new_units): - self.sunits = self._convert_string_unit(new_units) + converted_units = self._convert_string_unit(new_units) + old_sunits = self.sunits + self.sunits = converted_units qsunits = self._get_qs_unit() - if new_units == "degC" and str(self.data_with_units.units) == "1.0 K": - self.data_with_units = qs.Quantity( - self.data_with_units.magnitude[:] - 273.15, new_units - ) - else: - self.data_with_units = self.data_with_units.rescale(qsunits) - self.raw_data_with_units = self.raw_data_with_units.rescale(qsunits) + try: + if new_units == "degC" and str(self.data_with_units.units) == "1.0 K": + self.data_with_units = qs.Quantity( + self.data_with_units.magnitude[:] - 273.15, new_units + ) + else: + self.data_with_units = self.data_with_units.rescale(qsunits) + self.raw_data_with_units = self.raw_data_with_units.rescale(qsunits) + except (ValueError, LookupError) as e: + self.sunits = old_sunits + self._raise_unit_error(e, "convert units", requested_units=new_units) self.data = self.data_with_units.magnitude[:] self.raw_data = self.raw_data_with_units.magnitude self.set_label() @@ -349,7 +382,10 @@ def _arithmetic_op(self, other, op, symbol): "Arithmetic operations require a PsData object or numeric " "scalar, got {}".format(type(other)) ) - result_quantity = op(self.data_with_units, other_val) + try: + result_quantity = op(self.data_with_units, other_val) + except (ValueError, LookupError) as e: + self._raise_unit_error(e, "arithmetic '{}'".format(symbol), other=other) result_key = "({} {} {})".format(self.data_key, symbol, other_label) return PsData( data_key=result_key, @@ -371,7 +407,12 @@ def _r_arithmetic_op(self, other, op, symbol): "Arithmetic operations require a PsData object or numeric " "scalar, got {}".format(type(other)) ) - result_quantity = op(other_val, self.data_with_units) + try: + result_quantity = op(other_val, self.data_with_units) + except (ValueError, LookupError) as e: + self._raise_unit_error( + e, "reflected arithmetic '{}'".format(symbol), other=other + ) result_key = "({} {} {})".format(other_label, symbol, self.data_key) return PsData( data_key=result_key, @@ -426,7 +467,10 @@ def __rpow__(self, other): ) def __neg__(self): - result_quantity = -1 * self.data_with_units + try: + result_quantity = -1 * self.data_with_units + except (ValueError, LookupError) as e: + self._raise_unit_error(e, "negation") result_key = "(-{})".format(self.data_key) return PsData( data_key=result_key, diff --git a/src/psPlotKit/data_manager/ps_data_exporter.py b/src/psPlotKit/data_manager/ps_data_exporter.py index f94321a..3f74864 100644 --- a/src/psPlotKit/data_manager/ps_data_exporter.py +++ b/src/psPlotKit/data_manager/ps_data_exporter.py @@ -1,4 +1,5 @@ import os +import re import csv import numpy as np from psPlotKit.util.logger import define_logger @@ -18,6 +19,22 @@ class PsDataExporter: save_location : str File path (single directory) or folder path (multiple directories) where the CSV output will be written. + first_key : str, optional + Data key that should always appear as the first column in the exported + CSV files. If provided and the key exists in the data, it will be + moved to the front of the column order. + export_keys : list, optional + List of keys to include in export. Each key is matched against data keys: + - For string keys, a match occurs if the key is in the list + - For tuple keys, a match occurs if any element of the tuple is in the list + If None, all non-internal keys are exported. + exact_keys : list, optional + List of keys to include in export with exact matching. The entire key + (including all tuple elements) must match exactly. If specified, this + takes precedence over export_keys. + skip_zero_data : bool, optional + If True, columns whose data are all zeros are excluded from export. + Defaults to False so existing behavior is preserved. Notes ----- @@ -28,12 +45,31 @@ class PsDataExporter: CSV file inside that folder. * Column headers are built from each :class:`PsData` object's ``data_label`` and ``mpl_units`` attributes (set by - :meth:`PsData.set_label`). + :meth:`PsData.set_label`). LaTeX/math-mode markers are removed + from ``mpl_units`` so that headers contain plain text (e.g. + ``$/m^3`` instead of ``$\\$$/$m^3$``). + * Internal keys (e.g., ``_zero_sentinel``) are automatically excluded + from export. + * For tuple data keys (e.g., ``("costing", "pump")``), multi-row headers + are generated with each tuple element in a separate row, making the + structure clearer when reading CSVs in spreadsheet applications. """ - def __init__(self, ps_data_manager, save_location): + def __init__( + self, + ps_data_manager, + save_location, + first_key=None, + export_keys=None, + exact_keys=None, + skip_zero_data=False, + ): self.ps_data_manager = ps_data_manager self.save_location = save_location + self.first_key = first_key + self.export_keys = export_keys + self.exact_keys = exact_keys + self.skip_zero_data = skip_zero_data @staticmethod def _ensure_csv_extension(path): @@ -49,6 +85,37 @@ def _strip_csv_extension(path): return path[:-4] return path + @staticmethod + def _clean_unit_for_export(units): + """Remove matplotlib/LaTeX math formatting from a unit string. + + :class:`PsData.set_label` wraps unit tokens in math mode for + matplotlib rendering (e.g. ``$\\$$`` for a dollar sign, + ``$m^3$`` for an exponent). CSV exports should contain plain + text such as ``$/m^3`` instead. + + Parameters + ---------- + units : str + The unit string as stored in ``PsData.mpl_units``. + + Returns + ------- + str + Unit string with LaTeX/math-mode markers removed. + """ + if not units or units == "-": + return units + + # Unwrap math-mode exponents first, e.g. $m^3$ -> m^3 + units = re.sub(r"\$([^\$]*\^[^\$]*)\$", r"\1", units) + # USD/kUSD/MUSD math-mode dollar signs: $\$$ -> $, k$\$$ -> k$, etc. + units = re.sub(r"([kM]?)\$\s*\\\$\s*\$", r"\1$", units) + # Drop any remaining LaTeX escape backslashes + units = units.replace("\\", "") + + return units + def export(self): """Export the data manager contents to CSV. @@ -79,6 +146,8 @@ def _group_data_by_directory(self): composite keys are plain strings) and multi-directory managers (where composite keys are tuples). + Internal keys like ``_zero_sentinel`` are excluded from export. + Returns ------- dict @@ -95,18 +164,158 @@ def _group_data_by_directory(self): dir_key = tuple(dir_key) data_key = dm._get_data_key(composite_key) + # Skip internal keys like _zero_sentinel + if self._is_internal_key(data_key): + continue + + # Skip keys not in export filter + if not self._should_export_key(data_key): + continue + + # Skip all-zero columns when requested + if self.skip_zero_data and self._is_zero_data(ps_data): + _logger.info("Skipping '{}': all values are zero".format(data_key)) + continue + if dir_key not in grouped: grouped[dir_key] = [] grouped[dir_key].append((data_key, ps_data)) - return grouped + def _is_zero_data(self, ps_data): + """Return True if every value in *ps_data* is exactly zero. + + Empty arrays are considered all-zero. + + Parameters + ---------- + ps_data : PsData + The data object to inspect. + + Returns + ------- + bool + True if the data array is empty or all elements are 0. + """ + data = ps_data.data + if data is None: + return True + arr = np.asarray(data) + + if arr.size == 0: + return True + arr = np.nan_to_num(arr, nan=0) + if np.sum(np.abs(arr)) < np.finfo(float).eps * 10 or np.isnan( + np.sum(np.abs(arr)) + ): + return True + return False + # return np.all(arr == 0 + + def _should_export_key(self, data_key): + """Check if a data key should be exported based on export_keys/exact_keys filters. + + Parameters + ---------- + data_key : str or tuple + The data key to check (may contain nested tuples). + + Returns + ------- + bool + True if the key should be included in export. + """ + # Flatten nested tuples for matching + flat_key = ( + self._flatten_key(data_key) if isinstance(data_key, tuple) else data_key + ) + + if self.exact_keys is not None: + return data_key in self.exact_keys + + if self.export_keys is not None: + if isinstance(flat_key, tuple): + # For tuple keys, match if ANY element is in export_keys + return any(elem in self.export_keys for elem in flat_key) + else: + return flat_key in self.export_keys + + return True + + def _is_internal_key(self, data_key): + """Check if a data key is an internal (non-exportable) key. + + Parameters + ---------- + data_key : str or tuple + The data key to check (may contain nested tuples). + + Returns + ------- + bool + True if the key should be excluded from export. + """ + + def check_item(item): + if isinstance(item, tuple): + return any("_zero_sentinel" in str(part) for part in item) + return "_zero_sentinel" in str(item) + + if isinstance(data_key, tuple): + # Check all elements including nested ones + for elem in data_key: + if check_item(elem): + return True + return False + return check_item(data_key) + + def _reorder_with_first(self, data_items): + """Reorder data items to place first_key at the front if specified. + + Parameters + ---------- + data_items : list[tuple] + List of (data_key, PsData) pairs. + + Returns + ------- + list[tuple] + Reordered list with first_key at the front if applicable. + """ + if not self.first_key: + return data_items + + reordered = [] + first_item = None + for item in data_items: + data_key = item[0] + # For tuple keys, also check flattened version for matching + key_to_check = ( + self._flatten_key(data_key) if isinstance(data_key, tuple) else data_key + ) + if data_key == self.first_key or ( + isinstance(key_to_check, tuple) and self.first_key in key_to_check + ): + first_item = item + else: + reordered.append(item) + + if first_item: + reordered.insert(0, first_item) + _logger.info("Placed '{}' as first column".format(self.first_key)) + else: + _logger.warning( + "first_key '{}' not found in data keys".format(self.first_key) + ) + + return reordered + def _build_header(self, data_items): - """Build a CSV header row from a list of (data_key, PsData) pairs. + """Build CSV header rows from a list of (data_key, PsData) pairs. - Each column header is formatted as ``"label (units)"`` using the - ``data_label`` and ``mpl_units`` attributes produced by - :meth:`PsData.set_label`. + For string data keys, returns a single header row. For tuple data keys, + returns multiple header rows with the tuple components in separate rows, + making multi-level keys easier to read. Parameters ---------- @@ -115,19 +324,136 @@ def _build_header(self, data_items): Returns ------- - list[str] - Column header strings. + list[list[str]] + List of header rows (each row is a list of strings). + Single row if all keys are strings, multiple rows for tuple keys. + """ + has_tuple_keys = any(isinstance(key, tuple) for key, _ in data_items) + + if not has_tuple_keys: + return self._build_single_row_header(data_items) + + return self._build_multi_row_header(data_items) + + def _build_single_row_header(self, data_items): + """Build a single header row from data items. + + Parameters + ---------- + data_items : list[tuple] + List of (data_key, PsData) pairs. + + Returns + ------- + list[list[str]] + Single row as a list containing one list of header strings. """ headers = [] for _, ps_data in data_items: label = ps_data.data_label - units = getattr(ps_data, "mpl_units", "-") + units = self._clean_unit_for_export(getattr(ps_data, "mpl_units", "-")) if units and units != "-": header = "{} ({})".format(label, units) else: header = str(label) headers.append(header) - return headers + return [headers] + + @staticmethod + def _flatten_key(data_key): + """Flatten nested tuples into a single-level tuple. + + For example, ("costing", ("stage 1", "pump"), "LCOW") becomes + ("costing", "stage 1", "pump", "LCOW"). + + Parameters + ---------- + data_key : tuple + The data key potentially containing nested tuples. + + Returns + ------- + tuple + Flattened tuple with all elements at single level. + """ + flattened = [] + for elem in data_key: + if isinstance(elem, (tuple, list)): + flattened.extend(elem) + else: + flattened.append(elem) + return tuple(flattened) + + def _build_multi_row_header(self, data_items): + """Build multi-row headers for tuple data keys. + + Parameters + ---------- + data_items : list[tuple] + List of (data_key, PsData) pairs. + + Returns + ------- + list[list[str]] + List of header rows. Shorter tuples are aligned to the bottom, + so all labels appear in the final row. + """ + # Flatten nested tuples to determine max depth + flattened_keys = [] + for key, _ in data_items: + if isinstance(key, (tuple, list)): + flattened_keys.append(self._flatten_key(key)) + else: + flattened_keys.append(key) + + max_depth = max( + (len(key) for key in flattened_keys if isinstance(key, tuple)), default=1 + ) + + header_rows = [[] for _ in range(max_depth)] + + for idx, (data_key, ps_data) in enumerate(data_items): + label = ps_data.data_label + units = self._clean_unit_for_export(getattr(ps_data, "mpl_units", "-")) + + flat_key = flattened_keys[idx] + + # If label equals data_key, use the last tuple element as label + if label == data_key and isinstance(data_key, (tuple, list)): + label = flat_key[-1] + if isinstance(label, (tuple, list)): + if len(label) == 1: + label = str(label[0]) + else: + label = ", ".join(label) + if units and units != "-": + + full_label = "{} ({})".format(label, units) + else: + full_label = str(label) + if isinstance(flat_key, (tuple, list)): + key_len = len(flat_key) + # Calculate starting row (align shorter tuples to bottom) + start_row = max_depth - key_len + for row_idx in range(max_depth): + if row_idx < start_row: + header_rows[row_idx].append("") + elif row_idx == max_depth - 1: + # Last row gets the full label + header_rows[row_idx].append(full_label) + else: + # Middle rows get tuple elements starting from start_row + elem_idx = row_idx - start_row + header_rows[row_idx].append(str(flat_key[elem_idx])) + else: + # String keys go in the last row only + for row_idx in range(max_depth): + if row_idx == max_depth - 1: + header_rows[row_idx].append(full_label) + else: + header_rows[row_idx].append("") + + return header_rows def _build_rows(self, data_items): """Build row data from a list of (data_key, PsData) pairs. @@ -164,15 +490,15 @@ def _build_rows(self, data_items): rows.append(row) return rows - def _write_csv(self, file_path, headers, rows): + def _write_csv(self, file_path, header_rows, rows): """Write headers and rows to a CSV file. Parameters ---------- file_path : str Destination file path. - headers : list[str] - Column header strings. + header_rows : list[list[str]] + List of header row lists (supports multi-level headers). rows : list[list] Row-major data. """ @@ -182,7 +508,8 @@ def _write_csv(self, file_path, headers, rows): with open(file_path, "w", newline="") as f: writer = csv.writer(f) - writer.writerow(headers) + for header_row in header_rows: + writer.writerow(header_row) writer.writerows(rows) _logger.info("Saved CSV: {}".format(file_path)) @@ -208,14 +535,15 @@ def _export_single(self, grouped): _logger.info("Single directory detected, exporting to single CSV file") data_items = list(grouped.values())[0] - headers = self._build_header(data_items) + data_items = self._reorder_with_first(data_items) + header_rows = self._build_header(data_items) rows = self._build_rows(data_items) _logger.info( "Writing {} columns and {} rows to {}".format( - len(headers), len(rows), save_path + len(header_rows[0]), len(rows), save_path ) ) - self._write_csv(save_path, headers, rows) + self._write_csv(save_path, header_rows, rows) return [save_path] def _dir_key_to_filename(self, dir_key): @@ -273,14 +601,15 @@ def _export_multiple(self, grouped): for dir_key, data_items in grouped.items(): filename = self._dir_key_to_filename(dir_key) + ".csv" file_path = os.path.join(folder, filename) - headers = self._build_header(data_items) + data_items = self._reorder_with_first(data_items) + header_rows = self._build_header(data_items) rows = self._build_rows(data_items) _logger.info( "Directory '{}': writing {} columns and {} rows to {}".format( - dir_key, len(headers), len(rows), filename + dir_key, len(header_rows[0]), len(rows), filename ) ) - self._write_csv(file_path, headers, rows) + self._write_csv(file_path, header_rows, rows) written.append(file_path) _logger.info( diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py index a053ccf..886720d 100644 --- a/src/psPlotKit/data_manager/ps_data_manager.py +++ b/src/psPlotKit/data_manager/ps_data_manager.py @@ -465,7 +465,14 @@ def get_data_keys(self, selected_keys): data_keys = self.data_keys[:] return data_keys - def export_data_to_csv(self, save_location): + def export_data_to_csv( + self, + save_location, + first_key=None, + export_keys=None, + exact_keys=None, + skip_zero_data=False, + ): """Export all loaded data to CSV file(s). Convenience wrapper around :class:`PsDataExporter`. If this @@ -478,6 +485,17 @@ def export_data_to_csv(self, save_location): save_location : str File path for single-directory export (e.g. ``"results.csv"``) or folder path for multi-directory export. + first_key : str, optional + Data key that should always appear as the first column in the + exported CSV files. + export_keys : list, optional + List of keys to include in export. For tuple keys, any matching + element will be included. If None, all non-internal keys are exported. + exact_keys : list, optional + List of keys to include in export with exact matching. The entire + key must match exactly. Takes precedence over export_keys. + skip_zero_data : bool, optional + If True, columns whose data are all zeros are excluded from export. Returns ------- @@ -486,7 +504,9 @@ def export_data_to_csv(self, save_location): """ from psPlotKit.data_manager.ps_data_exporter import PsDataExporter - exporter = PsDataExporter(self, save_location) + exporter = PsDataExporter( + self, save_location, first_key, export_keys, exact_keys, skip_zero_data + ) return exporter.export() def display_loaded_contents(self): diff --git a/src/psPlotKit/data_manager/tests/test_ps_data.py b/src/psPlotKit/data_manager/tests/test_ps_data.py index b5e91c9..7a1e493 100644 --- a/src/psPlotKit/data_manager/tests/test_ps_data.py +++ b/src/psPlotKit/data_manager/tests/test_ps_data.py @@ -382,3 +382,56 @@ def test_scalar_sub_and_div(self): a = PsData("a", "test", [10.0, 20.0, 30.0]) result = (a - 5) / 2 np.testing.assert_array_almost_equal(result.data, [2.5, 7.5, 12.5]) + + +# ---------- unit error handling ---------- + + +class TestUnitErrors: + def test_to_units_incompatible_raises_with_metadata(self, meter_ps): + with pytest.raises(ValueError) as excinfo: + meter_ps.to_units("kg") + msg = str(excinfo.value) + assert "length" in msg + assert "test" in msg + assert "kg" in msg + assert "m" in msg + assert excinfo.value.__cause__ is not None + + def test_to_units_preserves_original_units_on_failure(self, meter_ps): + with pytest.raises(ValueError): + meter_ps.to_units("kg") + assert meter_ps.sunits == "m" + np.testing.assert_array_almost_equal(meter_ps.data, [10.0, 20.0, 30.0]) + + def test_constructor_units_incompatible_raises_with_metadata(self): + with pytest.raises(ValueError) as excinfo: + PsData("length", "test", [1.0, 2.0], import_units="m", units="kg") + msg = str(excinfo.value) + assert "length" in msg + assert "test" in msg + assert "kg" in msg + + def test_assign_units_invalid_string_raises_with_metadata(self, meter_ps): + with pytest.raises(LookupError) as excinfo: + meter_ps.assign_units("totally_bogus_unit") + msg = str(excinfo.value) + assert "length" in msg + assert "test" in msg + assert "totally_bogus_unit" in msg + + def test_invalid_unit_string_raises_with_metadata(self): + with pytest.raises(LookupError) as excinfo: + PsData("x", "test", [1.0], import_units="totally_bogus_unit") + msg = str(excinfo.value) + assert "x" in msg + assert "test" in msg + assert "totally_bogus_unit" in msg + + def test_arithmetic_incompatible_units_raises_with_metadata(self, meter_ps, usd_ps): + with pytest.raises(ValueError) as excinfo: + meter_ps + usd_ps + msg = str(excinfo.value) + assert "length" in msg + assert "cost" in msg + assert "test" in msg diff --git a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py index bd014d0..bc704ae 100644 --- a/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py +++ b/src/psPlotKit/data_manager/tests/test_ps_data_exporter.py @@ -268,6 +268,199 @@ def test_dimensionless_units_no_parentheses(self, tmp_path): assert "(" not in headers[0] +class TestUnitFormattingCleanup: + def test_latex_dollar_sign_exported_as_plain_dollar(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "LCOW", + PsData("LCOW", "test", [1.0, 2.0], import_units="USD/m**3"), + ) + save_path = str(tmp_path / "lcow.csv") + exporter = PsDataExporter(dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header = next(reader)[0] + + assert "LCOW" in header + assert "$/m^3" in header + assert "$\\$$" not in header + assert "${" not in header + + def test_latex_cleanup_for_kusd_and_musd(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "capex_k", + PsData("capex_k", "test", [1.0], import_units="kUSD/m**3"), + ) + dm.add_data( + "dir_a", + "capex_M", + PsData("capex_M", "test", [1.0], import_units="MUSD/m**3"), + ) + save_path = str(tmp_path / "capex.csv") + exporter = PsDataExporter(dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert "k$/m^3" in headers[0] + assert "M$/m^3" in headers[1] + + def test_plain_units_passthrough_unchanged(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "metric_1", + PsData("metric_1", "test", [1.0, 2.0, 3.0], import_units="m"), + ) + save_path = str(tmp_path / "plain.csv") + exporter = PsDataExporter(dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header = next(reader)[0] + + assert "metric_1 (m)" == header + + def test_cleanup_multi_row_header_for_tuple_keys(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + ("costing", "LCOW"), + PsData("LCOW", "test", [1.0], import_units="USD/m**3"), + ) + save_path = str(tmp_path / "multirow.csv") + exporter = PsDataExporter(dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(2)] + + assert header_rows[0] == ["costing"] + assert "LCOW" in header_rows[1][0] + assert "$/m^3" in header_rows[1][0] + assert "$\\$$" not in header_rows[1][0] + + +class TestSkipZeroData: + def test_zero_column_skipped_when_enabled(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "nonzero", + PsData("nonzero", "test", [1.0, 2.0, 3.0]), + ) + dm.add_data( + "dir_a", + "all_zero", + PsData("all_zero", "test", [0.0, 0.0, 0.0]), + ) + save_path = str(tmp_path / "skip_zero.csv") + exporter = PsDataExporter(dm, save_path, skip_zero_data=True) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert len(headers) == 1 + assert "nonzero" in headers[0] + assert "all_zero" not in headers[0] + + def test_zero_column_kept_when_disabled(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "nonzero", + PsData("nonzero", "test", [1.0, 2.0, 3.0]), + ) + dm.add_data( + "dir_a", + "all_zero", + PsData("all_zero", "test", [0.0, 0.0, 0.0]), + ) + save_path = str(tmp_path / "keep_zero.csv") + exporter = PsDataExporter(dm, save_path, skip_zero_data=False) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert len(headers) == 2 + assert "all_zero" in headers[1] + + def test_single_zero_value_skipped(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "scalar_zero", + PsData("scalar_zero", "test", [0.0]), + ) + save_path = str(tmp_path / "scalar_zero.csv") + exporter = PsDataExporter(dm, save_path, skip_zero_data=True) + written = exporter.export() + + assert written == [] + assert not os.path.exists(save_path) + + def test_mixed_zero_and_nonzero_keeps_nonzero(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "mixed", + PsData("mixed", "test", [0.0, 1.0, 0.0]), + ) + save_path = str(tmp_path / "mixed.csv") + exporter = PsDataExporter(dm, save_path, skip_zero_data=True) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert "mixed" in headers[0] + + def test_skip_zero_data_via_manager_method(self, tmp_path): + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "nonzero", + PsData("nonzero", "test", [1.0, 2.0]), + ) + dm.add_data( + "dir_a", + "all_zero", + PsData("all_zero", "test", [0.0, 0.0]), + ) + save_path = str(tmp_path / "manager_skip_zero.csv") + written = dm.export_data_to_csv(save_path, skip_zero_data=True) + + with open(written[0], "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert len(headers) == 1 + assert "nonzero" in headers[0] + + # ---------- PsDataManager.export_data_to_csv convenience method ---------- @@ -352,3 +545,242 @@ def test_manual_single_via_manager(self, manual_single_dm, tmp_path): assert len(headers) == 2 assert len(rows) == 3 assert float(rows[0][0]) == pytest.approx(1.0) + + +class TestFirstKeyParameter: + def test_first_key_places_column_first_single_dir(self, manual_single_dm, tmp_path): + save_path = str(tmp_path / "first_key.csv") + exporter = PsDataExporter(manual_single_dm, save_path, first_key="metric_2") + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert "metric_2" in headers[0] + assert "metric_1" in headers[1] + + def test_first_key_places_column_first_multi_dir(self, manual_multi_dm, tmp_path): + save_folder = str(tmp_path / "first_key_export") + exporter = PsDataExporter(manual_multi_dm, save_folder, first_key="metric_2") + written = exporter.export() + + for path in written: + with open(path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + assert "metric_2" in headers[0] + + def test_first_key_not_found_logs_warning(self, manual_single_dm, tmp_path): + save_path = str(tmp_path / "missing_key.csv") + exporter = PsDataExporter(manual_single_dm, save_path, first_key="nonexistent") + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert len(headers) == 2 + assert "metric_1" in headers[0] + + def test_first_key_via_manager_method(self, manual_single_dm, tmp_path): + save_path = str(tmp_path / "manager_first_key.csv") + written = manual_single_dm.export_data_to_csv(save_path, first_key="metric_2") + + with open(written[0], "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert "metric_2" in headers[0] + + +class TestNestedTupleProcessing: + def test_nested_tuple_keys_flattened_in_header(self, tmp_path): + """Nested tuple keys are flattened for multi-row headers.""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + ("costing", ("stage 1", "pump"), "LCOW"), + PsData(("costing", ("stage 1", "pump"), "LCOW"), "test", [150.0]), + ) + dm.add_data( + "dir_a", + "simple_key", + PsData("Simple", "test", [50.0]), + ) + + save_path = str(tmp_path / "nested.csv") + exporter = PsDataExporter(dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(4)] + + assert len(header_rows) == 4 + assert header_rows[0] == ["costing", ""] + assert header_rows[1] == ["stage 1", ""] + assert header_rows[2] == ["pump", ""] + assert "LCOW" in header_rows[3][0] and "Simple" in header_rows[3][1] + + def test_nested_tuple_key_matching_in_export_keys(self, tmp_path): + """export_keys matches elements within nested tuples.""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + ("costing", ("stage 1", "pump"), "LCOW"), + PsData("pump LCOW", "test", [150.0]), + ) + dm.add_data( + "dir_a", + ("data", ("stage 2", "other"), "value"), + PsData("other value", "test", [200.0]), + ) + + save_path = str(tmp_path / "nested_filter.csv") + exporter = PsDataExporter(dm, save_path, export_keys=["stage 1"]) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(4)] + + assert len(header_rows[3]) == 1 + assert "LCOW" in header_rows[3][0] + + +class TestKeyFiltering: + def test_export_keys_filters_by_partial_match(self, tmp_path): + """export_keys filters keys by partial match (tuple elements).""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "water_recovery", + PsData("Water recovery", "test", [0.85]), + ) + dm.add_data( + "dir_a", + ("costing", "pump"), + PsData("pump", "test", [100.0]), + ) + dm.add_data( + "dir_a", + ("data_a", "data_b"), + PsData("data_b", "test", [200.0]), + ) + + save_path = str(tmp_path / "filtered.csv") + exporter = PsDataExporter( + dm, save_path, export_keys=["water_recovery", "costing"] + ) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(2)] + + assert len(header_rows[1]) == 2 + assert "Water recovery" in header_rows[1][0] + assert "pump" in header_rows[1][1] + + def test_exact_keys_filters_by_exact_match(self, tmp_path): + """exact_keys filters keys by exact match only.""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "water_recovery", + PsData("Water recovery", "test", [0.85]), + ) + dm.add_data( + "dir_a", + ("costing", "pump"), + PsData("pump", "test", [100.0]), + ) + dm.add_data( + "dir_a", + ("data_a", "data_b"), + PsData("data_b", "test", [200.0]), + ) + + save_path = str(tmp_path / "exact.csv") + exporter = PsDataExporter( + dm, save_path, exact_keys=["water_recovery", ("costing", "pump")] + ) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(2)] + + assert len(header_rows[1]) == 2 + assert "Water recovery" in header_rows[1][0] + assert "pump" in header_rows[1][1] + + def test_exact_keys_takes_precedence_over_export_keys(self, tmp_path): + """exact_keys takes precedence over export_keys.""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + "key_a", + PsData("Key A", "test", [1.0]), + ) + dm.add_data( + "dir_a", + ("costing", "pump"), + PsData("pump", "test", [100.0]), + ) + + save_path = str(tmp_path / "precedence.csv") + exporter = PsDataExporter( + dm, save_path, export_keys=["key_a"], exact_keys=[("costing", "pump")] + ) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(2)] + + assert len(header_rows[1]) == 1 + assert "pump" in header_rows[1][0] + + def test_no_filter_exports_all_non_internal_keys(self, manual_single_dm, tmp_path): + """Without filter, all non-internal keys are exported.""" + save_path = str(tmp_path / "no_filter.csv") + exporter = PsDataExporter(manual_single_dm, save_path) + exporter.export() + + with open(save_path, "r") as f: + reader = csv.reader(f) + headers = next(reader) + + assert len(headers) == 2 + + def test_export_keys_via_manager_method(self, tmp_path): + """export_data_to_csv on PsDataManager accepts export_keys parameter.""" + dm = PsDataManager() + dm.PsDataImportInstances = [] + dm.add_data( + "dir_a", + ("costing", "pump"), + PsData("pump", "test", [100.0]), + ) + dm.add_data( + "dir_a", + ("data_a", "data_b"), + PsData("data_b", "test", [200.0]), + ) + + save_path = str(tmp_path / "via_manager.csv") + dm.export_data_to_csv(save_path, export_keys=["costing"]) + + with open(save_path, "r") as f: + reader = csv.reader(f) + header_rows = [next(reader) for _ in range(2)] + + assert len(header_rows[1]) == 1 + assert "pump" in header_rows[1][0]