fix(F-09): address review-gate findings — resource leak + idempotency
- mathpix.py + figures.py: wrap fitz page-loop in try/finally so doc.close() is guaranteed even on exception (no file-handle leak) - ingest.py: DELETE FROM formulas/figures WHERE paper_id before re-insert so --rich re-ingest is idempotent (BIGSERIAL has no natural UNIQUE key; ON CONFLICT DO NOTHING was a no-op) Gate: 158 passed, ruff clean, mypy clean Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
@@ -220,13 +220,18 @@ def ingest_paper(
|
||||
|
||||
if formulas or figures:
|
||||
with get_conn() as conn2:
|
||||
with conn2.cursor() as cur:
|
||||
# Delete-before-insert keeps re-ingest idempotent (BIGSERIAL has no
|
||||
# natural UNIQUE key, so ON CONFLICT DO NOTHING never fires).
|
||||
cur.execute("DELETE FROM formulas WHERE paper_id = %s", (paper.id,))
|
||||
cur.execute("DELETE FROM figures WHERE paper_id = %s", (paper.id,))
|
||||
|
||||
if formulas:
|
||||
with conn2.cursor() as cur:
|
||||
cur.executemany(
|
||||
"""
|
||||
INSERT INTO formulas (paper_id, page, raw_latex, context, eq_label)
|
||||
VALUES (%s, %s, %s, %s, %s)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(f.paper_id, f.page, f.raw_latex, f.context, f.eq_label)
|
||||
@@ -241,7 +246,6 @@ def ingest_paper(
|
||||
"""
|
||||
INSERT INTO figures (paper_id, page, caption, image_path)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(fig.paper_id, fig.page, fig.caption, fig.image_path)
|
||||
|
||||
@@ -126,41 +126,44 @@ def extract_figures(
|
||||
logger.exception("fitz.open failed for %s", pdf_path)
|
||||
return results
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
images: list[Any] = page.get_images(full=True)
|
||||
text_blocks: list[Any] = page.get_text("blocks")
|
||||
try:
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
images: list[Any] = page.get_images(full=True)
|
||||
text_blocks: list[Any] = page.get_text("blocks")
|
||||
|
||||
for fig_idx, img_info in enumerate(images):
|
||||
xref: int = img_info[0]
|
||||
for fig_idx, img_info in enumerate(images):
|
||||
xref: int = img_info[0]
|
||||
|
||||
# Get image bounding box on the page
|
||||
img_rects: list[Any] = page.get_image_rects(xref)
|
||||
if not img_rects:
|
||||
continue
|
||||
img_rect = img_rects[0]
|
||||
# Get image bounding box on the page
|
||||
img_rects: list[Any] = page.get_image_rects(xref)
|
||||
if not img_rects:
|
||||
continue
|
||||
img_rect = img_rects[0]
|
||||
|
||||
# Extract and save the image pixmap
|
||||
try:
|
||||
pix: Any = fitz.Pixmap(doc, xref)
|
||||
if pix.n > 4: # CMYK or similar — convert to RGB
|
||||
pix = fitz.Pixmap(fitz.csRGB, pix)
|
||||
filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
|
||||
img_path = out_dir / filename
|
||||
pix.save(str(img_path))
|
||||
except Exception:
|
||||
logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
|
||||
continue
|
||||
# Extract and save the image pixmap
|
||||
try:
|
||||
pix: Any = fitz.Pixmap(doc, xref)
|
||||
if pix.n > 4: # CMYK or similar — convert to RGB
|
||||
pix = fitz.Pixmap(fitz.csRGB, pix)
|
||||
filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
|
||||
img_path = out_dir / filename
|
||||
pix.save(str(img_path))
|
||||
except Exception:
|
||||
logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
|
||||
continue
|
||||
|
||||
caption = _find_caption(img_rect, text_blocks)
|
||||
caption = _find_caption(img_rect, text_blocks)
|
||||
|
||||
results.append(
|
||||
FigureChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
image_path=str(img_path),
|
||||
caption=caption,
|
||||
results.append(
|
||||
FigureChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
image_path=str(img_path),
|
||||
caption=caption,
|
||||
)
|
||||
)
|
||||
)
|
||||
finally:
|
||||
doc.close()
|
||||
|
||||
return results
|
||||
|
||||
@@ -112,54 +112,60 @@ def _extract_formulas_pix2tex(pdf_path: str) -> list[FormulaChunk]:
|
||||
|
||||
scale = fitz.Matrix(3, 3)
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
blocks: list[Any] = page.get_text("blocks")
|
||||
try:
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
blocks: list[Any] = page.get_text("blocks")
|
||||
|
||||
for block in blocks:
|
||||
# block = (x0, y0, x1, y1, text, block_no, block_type)
|
||||
if len(block) < 5:
|
||||
continue
|
||||
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
|
||||
text: str = block[4]
|
||||
for block in blocks:
|
||||
# block = (x0, y0, x1, y1, text, block_no, block_type)
|
||||
if len(block) < 5:
|
||||
continue
|
||||
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
|
||||
text: str = block[4]
|
||||
|
||||
# Size filter
|
||||
width_pt = x1 - x0
|
||||
height_pt = y1 - y0
|
||||
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
|
||||
continue
|
||||
# Size filter
|
||||
width_pt = x1 - x0
|
||||
height_pt = y1 - y0
|
||||
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
|
||||
continue
|
||||
|
||||
# Math content filter
|
||||
if _math_symbol_count(text) < _MIN_MATH_SYMBOLS or _math_ratio(text) < _MIN_MATH_RATIO:
|
||||
continue
|
||||
# Math content filter
|
||||
if (
|
||||
_math_symbol_count(text) < _MIN_MATH_SYMBOLS
|
||||
or _math_ratio(text) < _MIN_MATH_RATIO
|
||||
):
|
||||
continue
|
||||
|
||||
# Crop at 3× for OCR quality
|
||||
clip_rect = fitz.Rect(x0, y0, x1, y1)
|
||||
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
|
||||
img = Image.frombytes(
|
||||
"RGB",
|
||||
(crop_pix.width, crop_pix.height),
|
||||
crop_pix.samples,
|
||||
)
|
||||
|
||||
try:
|
||||
latex: str = model(img)
|
||||
except Exception:
|
||||
logger.debug("pix2tex failed on page %d block, skipping", page_num + 1)
|
||||
continue
|
||||
|
||||
if not latex or not latex.strip():
|
||||
continue
|
||||
|
||||
context = text[:200].replace("\n", " ")
|
||||
results.append(
|
||||
FormulaChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
raw_latex=latex.strip(),
|
||||
context=context,
|
||||
# Crop at 3× for OCR quality
|
||||
clip_rect = fitz.Rect(x0, y0, x1, y1)
|
||||
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
|
||||
img = Image.frombytes(
|
||||
"RGB",
|
||||
(crop_pix.width, crop_pix.height),
|
||||
crop_pix.samples,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
latex: str = model(img)
|
||||
except Exception:
|
||||
logger.debug("pix2tex failed on page %d block, skipping", page_num + 1)
|
||||
continue
|
||||
|
||||
if not latex or not latex.strip():
|
||||
continue
|
||||
|
||||
context = text[:200].replace("\n", " ")
|
||||
results.append(
|
||||
FormulaChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
raw_latex=latex.strip(),
|
||||
context=context,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
doc.close()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
Reference in New Issue
Block a user