mainframe-migration-toolkit 0.2.0__py3-none-any.whl

mainframe_toolkit/sequential.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import csv
+import json
+from dataclasses import dataclass, field
+from decimal import Decimal
+from pathlib import Path
+from typing import Any, Iterable, Iterator, Mapping, Sequence
+
+from .cobol import cobol_decimal, decode_overpunch, decode_packed_decimal
+from .errors import DataFormatError, OptionalDependencyError
+
+
+@dataclass(frozen=True)
+class FieldSpec:
+    name: str
+    start: int
+    length: int
+    kind: str = "text"
+    scale: int = 0
+    encoding: str | None = None
+    strip: bool = True
+    signed: bool = True
+    byteorder: str = "big"
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, Any]) -> "FieldSpec":
+        return cls(**dict(value))
+
+
+@dataclass(frozen=True)
+class RecordLayout:
+    fields: tuple[FieldSpec, ...] = field(default_factory=tuple)
+    record_length: int | None = None
+    encoding: str = "utf-8"
+
+    def __post_init__(self) -> None:
+        names = [item.name for item in self.fields]
+        if len(names) != len(set(names)):
+            raise ValueError("layout field names must be unique")
+        for item in self.fields:
+            if item.start < 1 or item.length < 1:
+                raise ValueError(f"invalid field bounds for {item.name}")
+        required = max((item.start - 1 + item.length for item in self.fields), default=0)
+        if self.record_length is not None and self.record_length < required:
+            raise ValueError("record_length is shorter than the final field")
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, Any]) -> "RecordLayout":
+        return cls(
+            fields=tuple(FieldSpec.from_dict(item) for item in value.get("fields", [])),
+            record_length=value.get("record_length"),
+            encoding=value.get("encoding", "utf-8"),
+        )
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "RecordLayout":
+        with Path(path).open(encoding="utf-8") as stream:
+            return cls.from_dict(json.load(stream))
+
+    def parse(self, record: bytes) -> dict[str, Any]:
+        if self.record_length is not None and len(record) != self.record_length:
+            raise DataFormatError(
+                f"expected record length {self.record_length}, received {len(record)}"
+            )
+        result: dict[str, Any] = {}
+        for spec in self.fields:
+            raw = record[spec.start - 1 : spec.start - 1 + spec.length]
+            if len(raw) != spec.length:
+                raise DataFormatError(f"record is too short for field {spec.name}")
+            result[spec.name] = _decode_field(raw, spec, self.encoding)
+        return result
+
+
+def _decode_field(raw: bytes, spec: FieldSpec, default_encoding: str) -> Any:
+    kind = spec.kind.lower().replace("-", "_")
+    if kind in {"packed", "packed_decimal", "comp_3", "comp3"}:
+        return decode_packed_decimal(raw, spec.scale)
+    if kind in {"binary", "binary_int", "comp", "comp_5", "comp5"}:
+        return int.from_bytes(raw, byteorder=spec.byteorder, signed=spec.signed)
+
+    try:
+        text = raw.decode(spec.encoding or default_encoding)
+    except UnicodeDecodeError as exc:
+        raise DataFormatError(f"cannot decode field {spec.name}") from exc
+    value = text.strip() if spec.strip else text
+    if kind in {"text", "string", "alphanumeric", "x"}:
+        return value
+    if kind in {"integer", "int", "display_integer"}:
+        return int(value or "0")
+    if kind in {"decimal", "display_decimal", "numeric"}:
+        return cobol_decimal(value or "0", spec.scale)
+    if kind in {"zoned", "zoned_decimal", "overpunch"}:
+        return decode_overpunch(value or "0", spec.scale)
+    if kind == "bytes":
+        return raw
+    raise DataFormatError(f"unsupported field kind {spec.kind!r}")
+
+
+def _source_files(source: str | Path) -> list[Path]:
+    path = Path(source)
+    if path.is_file():
+        return [path]
+    if path.is_dir():
+        return sorted(item for item in path.rglob("*") if item.is_file() and not item.name.startswith("."))
+    raise FileNotFoundError(path)
+
+
+def _records_from_file(
+    path: Path,
+    *,
+    mode: str,
+    record_length: int | None,
+    allow_short: bool,
+) -> Iterator[tuple[int, bytes]]:
+    with path.open("rb") as stream:
+        if mode == "line":
+            for number, raw in enumerate(stream, 1):
+                yield number, raw.rstrip(b"\r\n")
+            return
+        if mode != "fixed":
+            raise ValueError("mode must be 'line' or 'fixed'")
+        if not record_length:
+            raise ValueError("record_length is required in fixed mode")
+        number = 0
+        while True:
+            raw = stream.read(record_length)
+            if not raw:
+                break
+            number += 1
+            if len(raw) != record_length and not allow_short:
+                raise DataFormatError(
+                    f"{path}: trailing record has {len(raw)} of {record_length} bytes"
+                )
+            yield number, raw
+
+
+def iter_sequential(
+    source: str | Path,
+    layout: RecordLayout,
+    *,
+    mode: str = "line",
+    on_error: str = "raise",
+    allow_short: bool = False,
+) -> Iterator[dict[str, Any]]:
+    """Stream line-sequential or fixed-block records as dictionaries."""
+    if on_error not in {"raise", "skip", "null"}:
+        raise ValueError("on_error must be raise, skip, or null")
+    for path in _source_files(source):
+        for number, raw in _records_from_file(
+            path,
+            mode=mode,
+            record_length=layout.record_length,
+            allow_short=allow_short,
+        ):
+            try:
+                yield layout.parse(raw)
+            except (DataFormatError, ValueError) as exc:
+                if on_error == "raise":
+                    raise DataFormatError(f"{path}:{number}: {exc}") from exc
+                if on_error == "null":
+                    yield {item.name: None for item in layout.fields}
+
+
+def _json_safe(value: Any) -> Any:
+    if isinstance(value, Decimal):
+        return format(value, "f")
+    if isinstance(value, bytes):
+        return value.hex()
+    return value
+
+
+def _destination_file(destination: str | Path, extension: str) -> Path:
+    path = Path(destination)
+    if path.suffix.lower() == extension:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        return path
+    path.mkdir(parents=True, exist_ok=True)
+    return path / f"part-00000{extension}"
+
+
+def _write_csv(
+    rows: Iterable[Mapping[str, Any]],
+    fields: Sequence[str],
+    destination: Path,
+    delimiter: str,
+) -> int:
+    count = 0
+    with destination.open("w", encoding="utf-8", newline="") as stream:
+        writer = csv.DictWriter(stream, fieldnames=fields, delimiter=delimiter)
+        writer.writeheader()
+        for row in rows:
+            writer.writerow({name: _json_safe(row.get(name)) for name in fields})
+            count += 1
+    return count
+
+
+def _write_parquet(
+    rows: Iterable[Mapping[str, Any]],
+    fields: Sequence[str],
+    destination: Path,
+    batch_size: int,
+) -> int:
+    try:
+        import pyarrow as arrow
+        import pyarrow.parquet as parquet
+    except ImportError as exc:
+        raise OptionalDependencyError("install mainframe-migration-toolkit[parquet]") from exc
+
+    writer = None
+    batch: list[dict[str, Any]] = []
+    count = 0
+    try:
+        for row in rows:
+            batch.append({name: row.get(name) for name in fields})
+            if len(batch) >= batch_size:
+                table = arrow.Table.from_pylist(batch)
+                writer = writer or parquet.ParquetWriter(destination, table.schema)
+                writer.write_table(table)
+                count += len(batch)
+                batch.clear()
+        if batch:
+            table = arrow.Table.from_pylist(batch)
+            writer = writer or parquet.ParquetWriter(destination, table.schema)
+            writer.write_table(table)
+            count += len(batch)
+        if writer is None:
+            schema = arrow.schema([(name, arrow.string()) for name in fields])
+            writer = parquet.ParquetWriter(destination, schema)
+    finally:
+        if writer is not None:
+            writer.close()
+    return count
+
+
+def convert_sequential(
+    source: str | Path,
+    destination: str | Path,
+    layout: RecordLayout,
+    *,
+    output_format: str | None = None,
+    mode: str = "line",
+    delimiter: str = ",",
+    on_error: str = "raise",
+    batch_size: int = 10_000,
+) -> dict[str, Any]:
+    """Convert sequential records to one CSV or Parquet part file."""
+    selected = (output_format or Path(destination).suffix.lstrip(".") or "csv").lower()
+    extension = ".parquet" if selected == "parquet" else ".csv"
+    output = _destination_file(destination, extension)
+    rows = iter_sequential(source, layout, mode=mode, on_error=on_error)
+    fields = [item.name for item in layout.fields]
+    if selected == "csv":
+        count = _write_csv(rows, fields, output, delimiter)
+    elif selected == "parquet":
+        count = _write_parquet(rows, fields, output, batch_size)
+    else:
+        raise ValueError("output_format must be csv or parquet")
+    return {"source": str(source), "output": str(output), "format": selected, "rows": count}