89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
from __future__ import annotations
|
|
from pathlib import Path
|
|
from typing import Dict
|
|
import xml.etree.ElementTree as ET
|
|
import polars as pl
|
|
from tools.base import BaseTool
|
|
|
|
|
|
class OutputDataTool(BaseTool):
|
|
def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
|
|
df = inputs.get("Input", pl.DataFrame())
|
|
if self.config is None or df.is_empty():
|
|
return {}
|
|
|
|
file_el = self.config.find("File")
|
|
if file_el is None or not file_el.text:
|
|
return {}
|
|
|
|
raw_path = (file_el.text or "").strip()
|
|
fmt = int(file_el.attrib.get("FileFormat", "0"))
|
|
max_records_str = (file_el.attrib.get("MaxRecords") or "").strip()
|
|
max_records = int(max_records_str) if max_records_str else None
|
|
|
|
opts = self.config.find("FormatSpecificOptions")
|
|
if opts is None:
|
|
opts = ET.Element("x")
|
|
|
|
multi_el = self.config.find("MultiFile")
|
|
multi_file = (multi_el.attrib.get("value", "False") if multi_el is not None else "False") == "True"
|
|
multi_field = (self.config.findtext("MultiFileField") or "").strip()
|
|
multi_type = (self.config.findtext("MultiFileType") or "Suffix").strip()
|
|
keep_field = (self.config.findtext("KeepField") or "True").strip().lower() == "true"
|
|
|
|
out_path = self.ctx.resolve_output_path(raw_path)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if multi_file and multi_field and multi_field in df.columns:
|
|
for val in df[multi_field].unique().to_list():
|
|
part = df.filter(pl.col(multi_field) == val)
|
|
if not keep_field:
|
|
part = part.drop(multi_field)
|
|
part_path = self._multi_path(out_path, str(val), multi_type)
|
|
self._write(part, part_path, fmt, opts)
|
|
elif max_records:
|
|
chunk_num = 0
|
|
for i in range(0, len(df), max_records):
|
|
chunk = df.slice(i, max_records)
|
|
chunk_path = out_path if chunk_num == 0 else out_path.with_stem(
|
|
f"{out_path.stem}_{chunk_num}"
|
|
)
|
|
self._write(chunk, chunk_path, fmt, opts)
|
|
chunk_num += 1
|
|
else:
|
|
self._write(df, out_path, fmt, opts)
|
|
|
|
if self.ctx.verbose:
|
|
print(f"[Output] Wrote {len(df)} rows → {out_path}")
|
|
|
|
return {}
|
|
|
|
def _multi_path(self, base: Path, value: str, mode: str) -> Path:
|
|
safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in value)
|
|
if mode == "Suffix":
|
|
return base.with_stem(f"{base.stem}_{safe}")
|
|
return base.with_stem(f"{safe}_{base.stem}")
|
|
|
|
def _write(self, df: pl.DataFrame, path: Path, fmt: int, opts: ET.Element) -> None:
|
|
if fmt in (0, 6): # CSV
|
|
delim = opts.findtext("Delimeter") or opts.findtext("Delimiter") or ","
|
|
# Accept both HeaderRow and Headers attribute names
|
|
header_val = (
|
|
opts.findtext("HeaderRow")
|
|
or opts.findtext("Headers")
|
|
or "True"
|
|
)
|
|
header = header_val.lower() != "false"
|
|
df.write_csv(str(path), separator=delim, include_header=header)
|
|
elif fmt == 25: # Excel
|
|
df.write_excel(str(path))
|
|
elif fmt == 2: # Parquet
|
|
df.write_parquet(str(path))
|
|
elif fmt == 19: # YXDB — fall back to Parquet
|
|
fallback = path.with_suffix(".parquet")
|
|
df.write_parquet(str(fallback))
|
|
if self.ctx.verbose:
|
|
print(f"[Output] YXDB write not supported; wrote Parquet to {fallback}")
|
|
else:
|
|
df.write_csv(str(path))
|