370 lines
14 KiB
Python
370 lines
14 KiB
Python
"""
|
|
Streaming reader for very large Amazon transaction .xlsx files.
|
|
|
|
Design goals (validated on 200-850 MB / 780k-row source files):
|
|
* Never load the whole worksheet or its styles into memory.
|
|
* Stream rows via xml.etree.iterparse over the worksheet part.
|
|
* Load sharedStrings once per file (needed to resolve text cells).
|
|
* Auto-detect the *raw data* worksheet and skip any user-made pivot/summary sheet
|
|
by locating the row whose headers contain settlement id + total + date/time.
|
|
|
|
Public API:
|
|
reader = TransactionReader(path)
|
|
reader.detect() -> fills sheet_name / header_row / column_mapping
|
|
for rec in reader.iter_records(): -> yields normalized dict records
|
|
reader.file_meta -> FileMeta (sheets, rows, dates, currency, marketplace)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import zipfile
|
|
from dataclasses import dataclass, field
|
|
from datetime import date
|
|
from typing import Iterator
|
|
from xml.etree import ElementTree as ET
|
|
|
|
from .column_map import (
|
|
AMOUNT_FIELDS, FIELD_KIND, ColumnMapping, build_mapping,
|
|
)
|
|
from .dates import parse_amazon_date_fast
|
|
|
|
NS = "{http://schemas.openxmlformats.org/spreadsheetml/2006/main}"
|
|
_R = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}"
|
|
|
|
_COL_RE = re.compile(r"^([A-Z]+)")
|
|
|
|
|
|
def _col_letters(ref: str | None) -> str | None:
|
|
if not ref:
|
|
return None
|
|
m = _COL_RE.match(ref)
|
|
return m.group(1) if m else None
|
|
|
|
|
|
def _col_to_index(letters: str) -> int:
|
|
n = 0
|
|
for ch in letters:
|
|
n = n * 26 + (ord(ch) - 64)
|
|
return n - 1
|
|
|
|
|
|
class ParseError(Exception):
|
|
"""Raised when a file is not a usable Amazon transaction report."""
|
|
|
|
|
|
@dataclass
|
|
class FileMeta:
|
|
filename: str = ""
|
|
size_bytes: int = 0
|
|
worksheets: list[str] = field(default_factory=list)
|
|
data_sheet: str = ""
|
|
header_row: int = 0
|
|
imported_rows: int = 0
|
|
min_date: date | None = None
|
|
max_date: date | None = None
|
|
currency: str = "USD"
|
|
marketplace: str | None = None
|
|
unmapped_headers: dict[str, str] = field(default_factory=dict)
|
|
missing_required: list[str] = field(default_factory=list)
|
|
# Finance-added translation/helper header rows found below the real header and skipped.
|
|
helper_rows_skipped: int = 0
|
|
# column-letter -> Σ of numeric values seen in columns with NO mapped field.
|
|
# Non-zero sums are surfaced as errors: no amount is ever silently excluded.
|
|
unmapped_amount_sums: dict[str, float] = field(default_factory=dict)
|
|
|
|
|
|
class TransactionReader:
|
|
# A sheet's header row must resolve at least these to be considered the data sheet.
|
|
_DETECT_REQUIRED = {"settlement_id", "total", "date_time"}
|
|
|
|
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._zip: zipfile.ZipFile | None = None
|
|
self._shared: list[str] | None = None
|
|
self._sheets: list[tuple[str, str]] = [] # (name, part)
|
|
self.sheet_name: str = ""
|
|
self.sheet_part: str = ""
|
|
self.header_row: int = 0
|
|
self.column_mapping: ColumnMapping | None = None
|
|
self.file_meta = FileMeta(filename=self.filename)
|
|
|
|
# -- lifecycle -----------------------------------------------------------
|
|
def open(self) -> None:
|
|
if self._zip is not None:
|
|
return
|
|
try:
|
|
self._zip = zipfile.ZipFile(self.path)
|
|
except zipfile.BadZipFile as e:
|
|
raise ParseError(
|
|
f"'{self.filename}' is not a valid .xlsx file (corrupt or truncated)."
|
|
) from e
|
|
self.file_meta.size_bytes = os.path.getsize(self.path)
|
|
|
|
def close(self) -> None:
|
|
if self._zip is not None:
|
|
self._zip.close()
|
|
self._zip = None
|
|
self._shared = None
|
|
|
|
def __enter__(self) -> "TransactionReader":
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, *exc) -> None:
|
|
self.close()
|
|
|
|
# -- workbook structure --------------------------------------------------
|
|
def _read_sheets(self) -> list[tuple[str, str]]:
|
|
if self._sheets:
|
|
return self._sheets
|
|
assert self._zip is not None
|
|
wb = self._zip.read("xl/workbook.xml").decode("utf-8", "replace")
|
|
rels = self._zip.read("xl/_rels/workbook.xml.rels").decode("utf-8", "replace")
|
|
# Attribute order varies by writer (Excel: Id..Target; openpyxl: Target..Id).
|
|
rid_to_target: dict[str, str] = {}
|
|
for rel in re.findall(r"<Relationship\b[^>]*/?>", rels):
|
|
idm = re.search(r'Id="([^"]+)"', rel)
|
|
tgt = re.search(r'Target="([^"]+)"', rel)
|
|
if idm and tgt:
|
|
rid_to_target[idm.group(1)] = tgt.group(1)
|
|
sheets: list[tuple[str, str]] = []
|
|
for s in re.findall(r"<sheet\b[^>]*/>", wb):
|
|
name = re.search(r'name="([^"]*)"', s)
|
|
rid = re.search(r'r:id="([^"]*)"', s)
|
|
if not name or not rid:
|
|
continue
|
|
target = rid_to_target.get(rid.group(1), "")
|
|
if not target:
|
|
continue
|
|
part = "xl/" + target if not target.startswith("/") else target.lstrip("/")
|
|
sheets.append((name.group(1), part))
|
|
self._sheets = sheets
|
|
self.file_meta.worksheets = [n for n, _ in sheets]
|
|
return sheets
|
|
|
|
def _shared_strings(self) -> list[str]:
|
|
if self._shared is not None:
|
|
return self._shared
|
|
assert self._zip is not None
|
|
out: list[str] = []
|
|
if "xl/sharedStrings.xml" in self._zip.namelist():
|
|
with self._zip.open("xl/sharedStrings.xml") as fh:
|
|
for _ev, el in ET.iterparse(fh, events=("end",)):
|
|
if el.tag == NS + "si":
|
|
out.append("".join(t.text or "" for t in el.iter(NS + "t")))
|
|
el.clear()
|
|
self._shared = out
|
|
return out
|
|
|
|
def _cell_text(self, c: ET.Element, shared: list[str]) -> str | None:
|
|
t = c.get("t")
|
|
v = c.find(NS + "v")
|
|
if t == "s":
|
|
if v is None or v.text is None:
|
|
return None
|
|
try:
|
|
return shared[int(v.text)]
|
|
except (ValueError, IndexError):
|
|
return None
|
|
if t == "inlineStr":
|
|
iss = c.find(NS + "is")
|
|
return "".join(x.text or "" for x in iss.iter(NS + "t")) if iss is not None else None
|
|
return v.text if v is not None else None
|
|
|
|
def _iter_sheet(
|
|
self, part: str, shared: list[str], limit: int | None = None
|
|
) -> Iterator[tuple[int, dict[str, str | None]]]:
|
|
"""Yield (row_number, {col_letter: resolved_text})."""
|
|
assert self._zip is not None
|
|
with self._zip.open(part) as fh:
|
|
n = 0
|
|
for _ev, el in ET.iterparse(fh, events=("end",)):
|
|
if el.tag != NS + "row":
|
|
continue
|
|
rnum = el.get("r")
|
|
cells: dict[str, str | None] = {}
|
|
for idx, c in enumerate(el):
|
|
if c.tag != NS + "c":
|
|
continue
|
|
col = _col_letters(c.get("r")) or _index_to_col(idx)
|
|
cells[col] = self._cell_text(c, shared)
|
|
try:
|
|
rn = int(rnum) if rnum else n + 1
|
|
except ValueError:
|
|
rn = n + 1
|
|
yield rn, cells
|
|
el.clear()
|
|
n += 1
|
|
if limit is not None and n >= limit:
|
|
return
|
|
|
|
# -- detection -----------------------------------------------------------
|
|
def detect(self) -> ColumnMapping:
|
|
"""Locate the raw data sheet + header row and build the column mapping."""
|
|
self.open()
|
|
shared = self._shared_strings()
|
|
best: tuple[int, str, str, int, ColumnMapping] | None = None # (score, name, part, hdr_row, mapping)
|
|
for name, part in self._read_sheets():
|
|
for rnum, cells in self._iter_sheet(part, shared, limit=15):
|
|
header_cells = [(col, txt) for col, txt in cells.items() if txt and str(txt).strip()]
|
|
if len(header_cells) < 5:
|
|
continue
|
|
mapping = build_mapping(header_cells, rnum, self.saved_overrides)
|
|
have = set(mapping.field_to_col)
|
|
if self._DETECT_REQUIRED.issubset(have):
|
|
score = len(have)
|
|
if best is None or score > best[0]:
|
|
best = (score, name, part, rnum, mapping)
|
|
break # header found on this sheet; move to next sheet
|
|
if best is None:
|
|
raise ParseError(
|
|
f"'{self.filename}': could not find an Amazon transaction header row "
|
|
f"(need columns: date/time, settlement id, total)."
|
|
)
|
|
_, name, part, hdr_row, mapping = best
|
|
self.sheet_name = name
|
|
self.sheet_part = part
|
|
self.header_row = hdr_row
|
|
self.column_mapping = mapping
|
|
self.file_meta.data_sheet = name
|
|
self.file_meta.header_row = hdr_row
|
|
self.file_meta.unmapped_headers = mapping.unmapped
|
|
self.file_meta.missing_required = mapping.missing_required
|
|
return mapping
|
|
|
|
# -- records -------------------------------------------------------------
|
|
def iter_records(self, only_fields: set[str] | None = None) -> Iterator[dict]:
|
|
"""
|
|
Stream normalized records from the detected data sheet.
|
|
Each record carries _source_file / _source_sheet / _source_row for traceability.
|
|
"""
|
|
if self.column_mapping is None:
|
|
self.detect()
|
|
assert self.column_mapping is not None and self._zip is not None
|
|
shared = self._shared_strings()
|
|
col_to_field = self.column_mapping.col_to_field
|
|
want = only_fields
|
|
min_d: date | None = None
|
|
max_d: date | None = None
|
|
marketplace: str | None = None
|
|
count = 0
|
|
# detect currency from preamble ("All amounts in USD, unless specified")
|
|
unmapped_sums = self.file_meta.unmapped_amount_sums
|
|
for rnum, cells in self._iter_sheet(self.sheet_part, shared):
|
|
if rnum <= self.header_row:
|
|
continue
|
|
rec: dict = {
|
|
"_source_file": self.filename,
|
|
"_source_sheet": self.sheet_name,
|
|
"_source_row": rnum,
|
|
}
|
|
has_value = False
|
|
for col, val in cells.items():
|
|
fld = col_to_field.get(col)
|
|
if not fld:
|
|
# No amount is silently excluded: sum numeric data in unmapped columns.
|
|
if val not in (None, ""):
|
|
try:
|
|
unmapped_sums[col] = unmapped_sums.get(col, 0.0) + float(val)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
continue
|
|
if want is not None and fld not in want:
|
|
continue
|
|
rec[fld] = _convert(fld, val)
|
|
if rec[fld] not in (None, ""):
|
|
has_value = True
|
|
if not has_value:
|
|
continue
|
|
# track meta cheaply
|
|
d = None
|
|
if "date_time" in rec and rec["date_time"]:
|
|
d = parse_amazon_date_fast(rec["date_time"])
|
|
rec["_date"] = d
|
|
# Localized EU reports carry a Finance-added translation header directly under
|
|
# the real header row — text where numbers belong. Skip it, count it.
|
|
if rnum <= self.header_row + 2 and d is None:
|
|
sid = str(rec.get("settlement_id") or "")
|
|
if not sid.replace(".", "").isdigit():
|
|
self.file_meta.helper_rows_skipped += 1
|
|
continue
|
|
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 quick_expected_rows(path: str) -> int:
|
|
"""
|
|
Fast (KB-sized) estimate of data-row count without loading sharedStrings, used to drive
|
|
the progress bar. Reads the _xlnm._FilterDatabase defined name (or the largest sheet's
|
|
<dimension>) to find the last row.
|
|
"""
|
|
try:
|
|
z = zipfile.ZipFile(path)
|
|
except Exception:
|
|
return 0
|
|
best = 0
|
|
try:
|
|
wb = z.read("xl/workbook.xml").decode("utf-8", "replace")
|
|
for end in re.findall(r"\$[A-Z]+\$\d+:\$[A-Z]+\$(\d+)", wb):
|
|
best = max(best, int(end) - 8) # data rows below the row-8 header
|
|
if best <= 0:
|
|
# fall back to the biggest worksheet's dimension
|
|
for name in z.namelist():
|
|
if name.startswith("xl/worksheets/sheet") and name.endswith(".xml"):
|
|
head = z.open(name).read(4096).decode("utf-8", "replace")
|
|
m = re.search(r'<dimension ref="[A-Z]+\d+:[A-Z]+(\d+)"', head)
|
|
if m:
|
|
best = max(best, int(m.group(1)) - 8)
|
|
except Exception:
|
|
return max(best, 0)
|
|
finally:
|
|
z.close()
|
|
return max(best, 0)
|
|
|
|
|
|
def _index_to_col(idx: int) -> str:
|
|
idx += 1
|
|
s = ""
|
|
while idx:
|
|
idx, r = divmod(idx - 1, 26)
|
|
s = chr(65 + r) + s
|
|
return s
|
|
|
|
|
|
def _convert(field_name: str, raw: str | None):
|
|
kind = FIELD_KIND.get(field_name, "text")
|
|
if kind == "amount":
|
|
if raw is None or raw == "":
|
|
return 0.0
|
|
try:
|
|
return float(raw)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|
|
if kind == "int":
|
|
if raw in (None, ""):
|
|
return None
|
|
try:
|
|
return int(float(raw))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
if raw is None:
|
|
return None
|
|
s = str(raw).strip()
|
|
if kind == "id" and s.endswith(".0"):
|
|
s = s[:-2]
|
|
return s or None
|