feat: Add new feature for CsvSplit per bytes.(fixes #691

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(

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
+```

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)