222 lines
8.6 KiB
Python
222 lines
8.6 KiB
Python
"""
|
|
Fast Amazon transaction reader backed by python-calamine (Rust).
|
|
|
|
~4-5x faster than the pure-Python iterparse reader for the large source files. It parses
|
|
a worksheet into memory once (higher peak RSS), so the streaming `TransactionReader` remains
|
|
the low-memory fallback (see `make_reader`). Public interface matches `TransactionReader`:
|
|
detect() / iter_records() / file_meta / column_mapping / close().
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from datetime import date
|
|
from typing import Iterator
|
|
|
|
from openpyxl.utils import get_column_letter
|
|
|
|
from .column_map import FIELD_KIND, ColumnMapping, build_mapping
|
|
from .dates import parse_amazon_date_fast
|
|
from .xlsx_reader import FileMeta, ParseError
|
|
|
|
_DETECT_REQUIRED = {"settlement_id", "total", "date_time"}
|
|
|
|
|
|
class CalamineReader:
|
|
def __init__(self, path: str, saved_overrides: dict[str, str] | None = None):
|
|
self.path = path
|
|
self.filename = os.path.basename(path)
|
|
self.saved_overrides = saved_overrides
|
|
self._wb = None
|
|
self._sheet = None # parsed CalamineSheet for the data sheet
|
|
self.sheet_name = ""
|
|
self.header_row = 0 # 1-based (Excel)
|
|
self.column_mapping: ColumnMapping | None = None
|
|
self._field_to_idx: dict[str, int] = {}
|
|
self.file_meta = FileMeta(filename=self.filename)
|
|
|
|
# -- lifecycle --
|
|
def open(self) -> None:
|
|
if self._wb is not None:
|
|
return
|
|
from python_calamine import CalamineWorkbook
|
|
try:
|
|
self._wb = CalamineWorkbook.from_path(self.path)
|
|
except Exception as e: # noqa: BLE001
|
|
raise ParseError(f"'{self.filename}' could not be opened as a spreadsheet: {e}") from e
|
|
self.file_meta.size_bytes = os.path.getsize(self.path)
|
|
self.file_meta.worksheets = list(self._wb.sheet_names)
|
|
|
|
def close(self) -> None:
|
|
self._wb = None
|
|
self._sheet = None
|
|
|
|
def __enter__(self):
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self.close()
|
|
|
|
# -- detection --
|
|
def detect(self) -> ColumnMapping:
|
|
self.open()
|
|
assert self._wb is not None
|
|
# Evaluate every sheet; pick the candidate with the most mapped fields (tie ->
|
|
# tallest sheet). This keeps user-made pivot/helper sheets from winning, and the
|
|
# top-down row scan means the real localized header (row 8) beats any translation
|
|
# helper row below it.
|
|
best = None
|
|
best_key = (-1, -1)
|
|
for name in self._wb.sheet_names:
|
|
sheet = self._wb.get_sheet_by_name(name)
|
|
head = sheet.to_python(nrows=15)
|
|
for r_idx, row in enumerate(head):
|
|
cells = [(get_column_letter(i + 1), str(v))
|
|
for i, v in enumerate(row) if v not in ("", None)]
|
|
if len(cells) < 5:
|
|
continue
|
|
mapping = build_mapping(cells, r_idx + 1, self.saved_overrides)
|
|
if _DETECT_REQUIRED.issubset(set(mapping.field_to_col)):
|
|
key = (len(mapping.field_to_col), self._safe_height(name))
|
|
if key > best_key:
|
|
best, best_key = (name, sheet, r_idx, mapping), key
|
|
break # first qualifying row per sheet (the real header)
|
|
if not best:
|
|
raise ParseError(
|
|
f"'{self.filename}': could not find an Amazon transaction header row "
|
|
f"(need columns: date/time, settlement id, total)."
|
|
)
|
|
name, sheet, r_idx, mapping = best
|
|
self.sheet_name = name
|
|
self._sheet = sheet
|
|
self.header_row = r_idx + 1
|
|
self.column_mapping = mapping
|
|
# field -> column index (0-based) from the letter mapping
|
|
self._field_to_idx = {
|
|
fld: _letter_to_idx(col) for col, fld in mapping.col_to_field.items()
|
|
}
|
|
self.file_meta.data_sheet = name
|
|
self.file_meta.header_row = self.header_row
|
|
self.file_meta.unmapped_headers = mapping.unmapped
|
|
self.file_meta.missing_required = mapping.missing_required
|
|
return mapping
|
|
|
|
def _safe_height(self, name: str) -> int:
|
|
try:
|
|
return self._wb.get_sheet_by_name(name).total_height # type: ignore[union-attr]
|
|
except Exception:
|
|
return 1 << 30
|
|
|
|
# -- records --
|
|
def iter_records(self, only_fields: set[str] | None = None) -> Iterator[dict]:
|
|
if self.column_mapping is None:
|
|
self.detect()
|
|
assert self._sheet is not None
|
|
idx_map = self._field_to_idx
|
|
if only_fields is not None:
|
|
idx_map = {f: i for f, i in idx_map.items() if f in only_fields}
|
|
items = list(idx_map.items())
|
|
mapped_idx = set(self._field_to_idx.values())
|
|
unmapped_sums = self.file_meta.unmapped_amount_sums
|
|
hdr = self.header_row # 1-based; data starts at hdr+1 (Excel) => row index hdr (0-based)
|
|
min_d: date | None = None
|
|
max_d: date | None = None
|
|
marketplace: str | None = None
|
|
count = 0
|
|
want_date = "date_time" in idx_map
|
|
excel_row = 0
|
|
for excel_row, row in enumerate(self._sheet.iter_rows(), start=1):
|
|
if excel_row <= hdr:
|
|
continue
|
|
rec = {"_source_file": self.filename, "_source_sheet": self.sheet_name,
|
|
"_source_row": excel_row}
|
|
has_value = False
|
|
for fld, i in items:
|
|
v = row[i] if i < len(row) else None
|
|
cv = _conv_cal(fld, v)
|
|
rec[fld] = cv
|
|
if cv not in (None, "", 0.0):
|
|
has_value = True
|
|
elif cv == 0.0:
|
|
has_value = True # a zero amount is still a real cell
|
|
if not has_value:
|
|
continue
|
|
# Localized EU reports have a Finance-added translation header right below the
|
|
# real header — text where the settlement id / date belong. Skip and count it.
|
|
if excel_row <= hdr + 2:
|
|
raw_sid = row[self._field_to_idx["settlement_id"]] \
|
|
if self._field_to_idx.get("settlement_id", 99999) < len(row) else None
|
|
if not isinstance(raw_sid, (int, float)):
|
|
d_probe = parse_amazon_date_fast(rec.get("date_time")) \
|
|
if rec.get("date_time") else None
|
|
if d_probe is None:
|
|
self.file_meta.helper_rows_skipped += 1
|
|
continue
|
|
# No amount silently excluded: sum numeric values in unmapped columns.
|
|
# (Skipped entirely when every column is mapped — the common USA case.)
|
|
if len(row) > len(mapped_idx):
|
|
for i, v in enumerate(row):
|
|
if i not in mapped_idx and isinstance(v, (int, float)) and v:
|
|
col = get_column_letter(i + 1)
|
|
unmapped_sums[col] = unmapped_sums.get(col, 0.0) + float(v)
|
|
if want_date and rec.get("date_time"):
|
|
d = parse_amazon_date_fast(rec["date_time"])
|
|
rec["_date"] = d
|
|
if d:
|
|
if min_d is None or d < min_d:
|
|
min_d = d
|
|
if max_d is None or d > max_d:
|
|
max_d = d
|
|
if marketplace is None and rec.get("marketplace"):
|
|
marketplace = rec["marketplace"]
|
|
count += 1
|
|
yield rec
|
|
self.file_meta.imported_rows = count
|
|
self.file_meta.min_date = min_d
|
|
self.file_meta.max_date = max_d
|
|
self.file_meta.marketplace = marketplace
|
|
|
|
|
|
def _letter_to_idx(letters: str) -> int:
|
|
n = 0
|
|
for ch in letters:
|
|
n = n * 26 + (ord(ch) - 64)
|
|
return n - 1
|
|
|
|
|
|
def _conv_cal(field_name: str, v):
|
|
kind = FIELD_KIND.get(field_name, "text")
|
|
if kind == "amount":
|
|
if isinstance(v, bool):
|
|
return 0.0
|
|
if isinstance(v, (int, float)):
|
|
return float(v)
|
|
if v in ("", None):
|
|
return 0.0
|
|
try:
|
|
return float(v)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
if kind == "int":
|
|
if isinstance(v, bool):
|
|
return None
|
|
if isinstance(v, (int, float)):
|
|
return int(v)
|
|
if v in ("", None):
|
|
return None
|
|
try:
|
|
return int(float(v))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
if v in (None, ""):
|
|
return None
|
|
if isinstance(v, float):
|
|
s = str(int(v)) if v.is_integer() else repr(v)
|
|
elif isinstance(v, int):
|
|
s = str(v)
|
|
else:
|
|
s = str(v).strip()
|
|
if kind == "id" and s.endswith(".0"):
|
|
s = s[:-2]
|
|
return s or None
|