diff --git a/cliboa/scenario/transform/csv.py b/cliboa/scenario/transform/csv.py index e41f8324..c42b5395 100644 --- a/cliboa/scenario/transform/csv.py +++ b/cliboa/scenario/transform/csv.py @@ -14,6 +14,7 @@ import csv import glob import hashlib +import io import os import re import shutil @@ -1079,11 +1080,18 @@ class CsvSplit(FileBaseTransform): """ class Arguments(FileBaseTransform.Arguments): - method: Literal["rows", "grouped"] + method: Literal["rows", "grouped", "bytes"] key_column: str | None = None rows: int | None = None + max_bytes: int | None = None suffix_format: str = ".{:02d}" + @model_validator(mode="after") + def validate_method_requirements(self) -> "CsvSplit.Arguments": + if self.method == "bytes" and self.max_bytes is None: + raise InvalidParameter("max_bytes is required when method is 'bytes'.") + return self + def execute(self, *args) -> None: files = self.get_src_files() self.check_file_existence(files) @@ -1092,6 +1100,8 @@ def execute(self, *args) -> None: executeInstance = _CsvSplitMethodRows(self.args) elif self.args.method == "grouped": executeInstance = _CsvSplitMethodGrouped(self.args) + elif self.args.method == "bytes": + executeInstance = _CsvSplitMethodBytes(self.args) else: raise NotImplementedError( f"Defined {self.args.method} is not implemented logic in execute." @@ -1171,6 +1181,65 @@ def _split_one(self, filepath: str) -> None: ) +class _CsvSplitMethodBytes(_CsvSplitMethodBase): + def execute(self, files: list[str]) -> None: + for filepath in files: + self._split_one(filepath) + + def _split_one(self, filepath: str) -> None: + self._logger.info("Split {:s} per {:d} bytes".format(filepath, self.args.max_bytes)) + file_name, ext = os.path.splitext(os.path.basename(filepath)) + with open(filepath, "r", encoding=self.args.encoding, newline="") as f_in: + reader = csv.reader(f_in) + try: + header = next(reader) + except StopIteration: + self._logger.error(f"Empty {filepath}") + return + + header_line = self._serialize_row(header) + header_bytes = len(header_line.encode(self.args.encoding)) + + file_index = 0 + f_out = None + output_filepath = None + written_bytes = 0 + + for row in reader: + line = self._serialize_row(row) + line_bytes = len(line.encode(self.args.encoding)) + if header_bytes + line_bytes > self.args.max_bytes: + self._logger.warning( + f"A record in {filepath} exceeds max_bytes={self.args.max_bytes}" + f" even in a file of its own" + f" (header {header_bytes} bytes + record {line_bytes} bytes)." + ) + if f_out is None or written_bytes + line_bytes > self.args.max_bytes: + if f_out: + f_out.close() + self._logger.info( + f"Generated {output_filepath} with {written_bytes} bytes." + ) + suffix = self.args.suffix_format.format(file_index) + output_filepath = os.path.join( + self.args.resolve_dest_dir(), f"{file_name}{suffix}{ext}" + ) + f_out = open(output_filepath, "w", encoding=self.args.encoding, newline="") + f_out.write(header_line) + written_bytes = header_bytes + file_index += 1 + f_out.write(line) + written_bytes += line_bytes + if f_out: + f_out.close() + self._logger.info(f"Generated {output_filepath} with {written_bytes} bytes.") + + def _serialize_row(self, row: list[str]) -> str: + buf = io.StringIO() + csv.writer(buf).writerow(row) + return buf.getvalue() + + class _CsvSplitMethodGrouped(_CsvSplitMethodBase): def execute(self, files: list[str]) -> None: valid2 = EssentialParameters( diff --git a/docs/modules/csv_split.md b/docs/modules/csv_split.md index ca105d65..fbf9ea4b 100644 --- a/docs/modules/csv_split.md +++ b/docs/modules/csv_split.md @@ -7,10 +7,11 @@ Split csv files by specified method. |src_dir|Path of the directory which input files are places.|Yes|None|| |src_pattern|File pattern of source csv files. Regexp is available.|Yes|None|| |dest_dir|Path of the directory which is for output files.|No|None|If a non-existent directory path is specified, the directory is automatically created.| -|method|Split method.|Yes|None|Only `rows` or `grouped` can be specified.| +|method|Split method.|Yes|None|Only `rows`, `grouped` or `bytes` can be specified.| |rows|When method is `rows`, split every N rows.|No|None|Required when method is `rows`| -|suffix_format|When method is `rows`, output file's suffix.(used in python's str.format)|No|None|| +|suffix_format|When method is `rows` or `bytes`, output file's suffix.(used in python's str.format)|No|None|| |key_column|When method is `grouped`, column name to use grouped split.|No|None|Required when method is `grouped`| +|max_bytes|When method is `bytes`, upper limit of each output file size in bytes.|No|None|Required when method is `bytes`| |encoding|Character encoding when read and write|No|utf-8|| # Examples @@ -82,3 +83,19 @@ Output: /out/C.csv name, class epsilon, C ``` + + +## Method: bytes +Each output file contains the same header as the source file, and its size (header included) does not exceed `max_bytes`. A record is never split across files. + +``` +scenario: +- step: Split file by byte size + class: CsvSplit + arguments: + src_dir: /in + src_pattern: data\.csv + dest_dir: /out + method: bytes + max_bytes: 10485115 +``` diff --git a/tests/scenario/transform/test_csv.py b/tests/scenario/transform/test_csv.py index 6d817840..ea2fe680 100644 --- a/tests/scenario/transform/test_csv.py +++ b/tests/scenario/transform/test_csv.py @@ -3041,3 +3041,210 @@ def test_execute_ng_multiple_different_column_file(self): with pytest.raises(Exception): instance.execute() + + +class TestCsvSplitBytes(TestCsvTransform): + def _assert_split_files(self, dir_path, expected_results, max_bytes): + csv_files = [v for v in os.listdir(dir_path) if v.endswith(".csv")] + for file_name, expected_data in expected_results.items(): + assert ( + file_name in csv_files + ), f"Expected output {file_name} was not found, only exists {csv_files}" + + file_path = os.path.join(dir_path, file_name) + actual_size = os.path.getsize(file_path) + assert actual_size <= max_bytes, ( + f"Assertion failed for {file_name}: " + f"File size {actual_size} exceeds max_bytes {max_bytes}." + ) + with open(file_path, newline="") as f: + reader = csv.reader(f) + actual_data = [row for row in reader] + assert actual_data == expected_data, ( + f"Assertion failed for {file_name}: Data mismatch.\n" + f"Expected: {expected_data}\nActual: {actual_data}" + ) + + def test_execute_ok(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "alpha"], + ["2", "beta"], + ["3", "gamma"], + ["4", "delta"], + ["5", "epsilon"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + "max_bytes": 30, + } + ) + instance.execute() + + expected_results = { + "test1.00.csv": [["no", "name"], ["1", "alpha"], ["2", "beta"]], + "test1.01.csv": [["no", "name"], ["3", "gamma"], ["4", "delta"]], + "test1.02.csv": [["no", "name"], ["5", "epsilon"]], + } + self._assert_split_files(self._result_dir, expected_results, 30) + + def test_execute_ok_with_custom_suffix(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "alpha"], + ["2", "beta"], + ["3", "gamma"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + "max_bytes": 30, + "suffix_format": "_{:03d}", + } + ) + instance.execute() + + expected_results = { + "test1_000.csv": [["no", "name"], ["1", "alpha"], ["2", "beta"]], + "test1_001.csv": [["no", "name"], ["3", "gamma"]], + } + self._assert_split_files(self._result_dir, expected_results, 30) + + def test_execute_ok_quoted_newline_kept_in_one_record(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "multi\nline"], + ["2", "beta"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + "max_bytes": 25, + } + ) + instance.execute() + + expected_results = { + "test1.00.csv": [["no", "name"], ["1", "multi\nline"]], + "test1.01.csv": [["no", "name"], ["2", "beta"]], + } + self._assert_split_files(self._result_dir, expected_results, 25) + + def test_execute_ok_size_counted_in_encoded_bytes(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "あいう"], + ["2", "え"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + "max_bytes": 25, + } + ) + instance.execute() + + expected_results = { + "test1.00.csv": [["no", "name"], ["1", "あいう"]], + "test1.01.csv": [["no", "name"], ["2", "え"]], + } + self._assert_split_files(self._result_dir, expected_results, 25) + + def test_execute_ok_oversized_record_warns_and_is_written(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "alpha"], + ["2", "beta"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + "max_bytes": 15, + } + ) + with self.assertLogs( + "cliboa.scenario.transform.csv._CsvSplitMethodBytes", level="WARNING" + ) as cm: + instance.execute() + assert any("exceeds max_bytes" in message for message in cm.output) + + csv_files = [v for v in os.listdir(self._result_dir) if v.endswith(".csv")] + expected_results = { + "test1.00.csv": [["no", "name"], ["1", "alpha"]], + "test1.01.csv": [["no", "name"], ["2", "beta"]], + } + for file_name, expected_data in expected_results.items(): + assert ( + file_name in csv_files + ), f"Expected output {file_name} was not found, only exists {csv_files}" + + with open(os.path.join(self._result_dir, file_name), newline="") as f: + reader = csv.reader(f) + actual_data = [row for row in reader] + assert actual_data == expected_data, ( + f"Assertion failed for {file_name}: Data mismatch.\n" + f"Expected: {expected_data}\nActual: {actual_data}" + ) + + def test_execute_ng_no_max_bytes(self): + # create test file + csv_list1 = [ + ["no", "name"], + ["1", "alpha"], + ] + self._create_csv(csv_list1, fname="test1.csv") + + with pytest.raises(InvalidParameter) as execinfo: + # set the essential attributes + instance = CsvSplit() + instance._set_arguments( + { + "src_dir": self._data_dir, + "src_pattern": r"test1\.csv", + "dest_dir": self._result_dir, + "method": "bytes", + } + ) + instance.execute() + assert "max_bytes is required when method is 'bytes'." in str(execinfo.value)