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:
@@ -220,13 +220,18 @@ def ingest_paper(
|
|||||||
|
|
||||||
if formulas or figures:
|
if formulas or figures:
|
||||||
with get_conn() as conn2:
|
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:
|
if formulas:
|
||||||
with conn2.cursor() as cur:
|
with conn2.cursor() as cur:
|
||||||
cur.executemany(
|
cur.executemany(
|
||||||
"""
|
"""
|
||||||
INSERT INTO formulas (paper_id, page, raw_latex, context, eq_label)
|
INSERT INTO formulas (paper_id, page, raw_latex, context, eq_label)
|
||||||
VALUES (%s, %s, %s, %s, %s)
|
VALUES (%s, %s, %s, %s, %s)
|
||||||
ON CONFLICT DO NOTHING
|
|
||||||
""",
|
""",
|
||||||
[
|
[
|
||||||
(f.paper_id, f.page, f.raw_latex, f.context, f.eq_label)
|
(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)
|
INSERT INTO figures (paper_id, page, caption, image_path)
|
||||||
VALUES (%s, %s, %s, %s)
|
VALUES (%s, %s, %s, %s)
|
||||||
ON CONFLICT DO NOTHING
|
|
||||||
""",
|
""",
|
||||||
[
|
[
|
||||||
(fig.paper_id, fig.page, fig.caption, fig.image_path)
|
(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)
|
logger.exception("fitz.open failed for %s", pdf_path)
|
||||||
return results
|
return results
|
||||||
|
|
||||||
for page_num in range(len(doc)):
|
try:
|
||||||
page = doc[page_num]
|
for page_num in range(len(doc)):
|
||||||
images: list[Any] = page.get_images(full=True)
|
page = doc[page_num]
|
||||||
text_blocks: list[Any] = page.get_text("blocks")
|
images: list[Any] = page.get_images(full=True)
|
||||||
|
text_blocks: list[Any] = page.get_text("blocks")
|
||||||
|
|
||||||
for fig_idx, img_info in enumerate(images):
|
for fig_idx, img_info in enumerate(images):
|
||||||
xref: int = img_info[0]
|
xref: int = img_info[0]
|
||||||
|
|
||||||
# Get image bounding box on the page
|
# Get image bounding box on the page
|
||||||
img_rects: list[Any] = page.get_image_rects(xref)
|
img_rects: list[Any] = page.get_image_rects(xref)
|
||||||
if not img_rects:
|
if not img_rects:
|
||||||
continue
|
continue
|
||||||
img_rect = img_rects[0]
|
img_rect = img_rects[0]
|
||||||
|
|
||||||
# Extract and save the image pixmap
|
# Extract and save the image pixmap
|
||||||
try:
|
try:
|
||||||
pix: Any = fitz.Pixmap(doc, xref)
|
pix: Any = fitz.Pixmap(doc, xref)
|
||||||
if pix.n > 4: # CMYK or similar — convert to RGB
|
if pix.n > 4: # CMYK or similar — convert to RGB
|
||||||
pix = fitz.Pixmap(fitz.csRGB, pix)
|
pix = fitz.Pixmap(fitz.csRGB, pix)
|
||||||
filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
|
filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
|
||||||
img_path = out_dir / filename
|
img_path = out_dir / filename
|
||||||
pix.save(str(img_path))
|
pix.save(str(img_path))
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
|
logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
caption = _find_caption(img_rect, text_blocks)
|
caption = _find_caption(img_rect, text_blocks)
|
||||||
|
|
||||||
results.append(
|
results.append(
|
||||||
FigureChunk(
|
FigureChunk(
|
||||||
paper_id=paper_id,
|
paper_id=paper_id,
|
||||||
page=page_num + 1,
|
page=page_num + 1,
|
||||||
image_path=str(img_path),
|
image_path=str(img_path),
|
||||||
caption=caption,
|
caption=caption,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
finally:
|
||||||
|
doc.close()
|
||||||
|
|
||||||
return results
|
return results
|
||||||
|
|||||||
@@ -112,54 +112,60 @@ def _extract_formulas_pix2tex(pdf_path: str) -> list[FormulaChunk]:
|
|||||||
|
|
||||||
scale = fitz.Matrix(3, 3)
|
scale = fitz.Matrix(3, 3)
|
||||||
|
|
||||||
for page_num in range(len(doc)):
|
try:
|
||||||
page = doc[page_num]
|
for page_num in range(len(doc)):
|
||||||
blocks: list[Any] = page.get_text("blocks")
|
page = doc[page_num]
|
||||||
|
blocks: list[Any] = page.get_text("blocks")
|
||||||
|
|
||||||
for block in blocks:
|
for block in blocks:
|
||||||
# block = (x0, y0, x1, y1, text, block_no, block_type)
|
# block = (x0, y0, x1, y1, text, block_no, block_type)
|
||||||
if len(block) < 5:
|
if len(block) < 5:
|
||||||
continue
|
continue
|
||||||
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
|
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
|
||||||
text: str = block[4]
|
text: str = block[4]
|
||||||
|
|
||||||
# Size filter
|
# Size filter
|
||||||
width_pt = x1 - x0
|
width_pt = x1 - x0
|
||||||
height_pt = y1 - y0
|
height_pt = y1 - y0
|
||||||
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
|
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Math content filter
|
# Math content filter
|
||||||
if _math_symbol_count(text) < _MIN_MATH_SYMBOLS or _math_ratio(text) < _MIN_MATH_RATIO:
|
if (
|
||||||
continue
|
_math_symbol_count(text) < _MIN_MATH_SYMBOLS
|
||||||
|
or _math_ratio(text) < _MIN_MATH_RATIO
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
|
||||||
# Crop at 3× for OCR quality
|
# Crop at 3× for OCR quality
|
||||||
clip_rect = fitz.Rect(x0, y0, x1, y1)
|
clip_rect = fitz.Rect(x0, y0, x1, y1)
|
||||||
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
|
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
|
||||||
img = Image.frombytes(
|
img = Image.frombytes(
|
||||||
"RGB",
|
"RGB",
|
||||||
(crop_pix.width, crop_pix.height),
|
(crop_pix.width, crop_pix.height),
|
||||||
crop_pix.samples,
|
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,
|
|
||||||
)
|
)
|
||||||
)
|
|
||||||
|
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
|
return results
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user