fix: CMYK image drop, empty-doc crash, page_index shadowing, sqlite hardening, flaky tests

Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272.

- pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without-
  alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png
  then raises "unsupported colorspace", silently dropped by the surrounding
  except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves
  RGBA untouched).

- pipeline.py: detect_strategy([]) (an empty/whitespace-only source file)
  returned "content_based", routing into the PDF-oriented TOC-detection
  pipeline -- wasting a real LLM call before raising IndexingError. Empty
  node lists now route to level_based, whose build_tree_from_levels([])
  returns an empty structure instantly with zero LLM calls.

- page_index.py (shim): pageindex/__init__.py binds the canonical `page_index`
  function as the package attribute, but this file is ALSO a real submodule
  of the same name -- importing it anywhere (import machinery, unconditional)
  overwrites that attribute with the module object, breaking
  `from pageindex import page_index; page_index(x)` for the rest of the
  process. Made the shim module itself callable (delegates to the real
  function via a ModuleType subclass), so whichever object ends up in that
  slot is callable regardless of import order.

- storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape
  on a duplicate name (new CollectionAlreadyExistsError); the collections
  table's CHECK constraint only validated the name's first character (GLOB
  '*' is a wildcard, not a regex quantifier over the preceding class) --
  fixed to validate the whole string, and SQLiteStorage now also validates in
  Python (it's a public StorageEngine usable directly, bypassing
  LocalBackend's own check).

- tests/test_review_fixes_2.py: two tests used a ContentNode with no `level`
  set, so build_index took the content_based path and made real (retried,
  slow, and -- with a valid key -- billable) LLM calls instead of testing the
  text-stripping logic they claimed to. Mocked out _content_based_pipeline.

- retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of
  the canonical parse_pages/get_pdf_page_content that had already drifted
  (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical
  now, so the legacy pageindex.get_page_content path can't silently regress
  again.

- parser/markdown.py: a leading UTF-8 BOM broke first-header detection
  (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only
  backtick fences were recognized as code blocks, so a '#'-prefixed line
  inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading --
  recognize both fence styles.

- run_pageindex.py: --if-thinning wasn't migrated to the bare-flag +
  legacy-yes/no convention the other four --if-add-* flags got; bare usage
  raised an argparse error and it never went through the shared coercion.

- types.py: DocumentDetail's `structure` field was inside the class's
  total=False body, so TypedDict rules made it optional even though every
  backend always populates it. Split into a required base class.

Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped
(one pre-existing, unrelated flaky cloud-streaming test).

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
mountain 2026-07-09 11:58:59 +08:00
parent b9d021916f
commit 4e6a13576d
18 changed files with 368 additions and 36 deletions

View file

@ -8,6 +8,11 @@ class CollectionNotFoundError(PageIndexError):
pass
class CollectionAlreadyExistsError(PageIndexError):
"""Collection already exists (create_collection, not get_or_create)."""
pass
class DocumentNotFoundError(PageIndexError):
"""Document ID not found."""
pass

View file

@ -5,6 +5,13 @@ from ..parser.protocol import ContentNode, ParsedDocument
def detect_strategy(nodes: list[ContentNode]) -> str:
"""Determine which indexing strategy to use based on node data."""
if not nodes:
# No content at all (e.g. an empty/whitespace-only source file) ->
# level_based's build_tree_from_levels([]) returns an empty structure
# immediately with zero LLM calls. content_based's TOC-detection
# pipeline needs real page content; on an empty page_list it wastes an
# LLM call and then still raises, for no benefit.
return "level_based"
if any(n.level is not None for n in nodes):
return "level_based"
return "content_based"

View file

@ -3,6 +3,8 @@
# pageindex/index/page_index.py (the single source of truth). This module
# re-exports it so legacy imports (`from pageindex.page_index import ...`,
# `from pageindex import page_index`) keep working.
import sys
import types
import warnings
warnings.warn(
@ -13,3 +15,26 @@ warnings.warn(
)
from .index.page_index import * # noqa: F401,F403,E402
# pageindex/__init__.py binds the FUNCTION `page_index` as the package
# attribute `pageindex.page_index` (`from .index.page_index import *`). But
# this file is ALSO a real submodule of the same name — the moment anything,
# anywhere in the process, does `import pageindex.page_index` (exactly what
# `from pageindex.page_index import X` triggers), Python's import machinery
# overwrites that package attribute with THIS module object, clobbering the
# function binding. Afterwards `from pageindex import page_index; page_index(x)`
# would raise "TypeError: 'module' object is not callable" — silently, and
# depending entirely on whether this submodule happened to be imported yet.
#
# Fix: make this module itself callable, delegating to the real function, so
# whichever object ends up sitting in the `pageindex.page_index` slot — the
# function or this module — is callable either way. Both `from pageindex.page_index
# import page_index_main` (module attribute access) and
# `from pageindex import page_index; page_index(x)` (call) keep working
# regardless of import order.
class _CallableModule(types.ModuleType):
def __call__(self, *args, **kwargs):
return page_index(*args, **kwargs)
sys.modules[__name__].__class__ = _CallableModule

View file

@ -12,7 +12,12 @@ class MarkdownParser:
path = Path(file_path)
model = kwargs.get("model")
with open(path, "r", encoding="utf-8") as f:
# utf-8-sig strips a leading BOM if present (common from Windows
# editors/exporters) and is otherwise identical to plain utf-8. Without
# it, a BOM-prefixed first line fails the header regex below (the BOM
# isn't whitespace, so .strip() doesn't remove it), misclassifying the
# document's first heading as unrecognized preamble text.
with open(path, "r", encoding="utf-8-sig") as f:
content = f.read()
lines = content.split("\n")
@ -23,7 +28,10 @@ class MarkdownParser:
def _extract_headers(self, lines: list[str]) -> list[dict]:
header_pattern = r"^(#{1,6})\s+(.+)$"
code_block_pattern = r"^```"
# CommonMark allows both backtick and tilde fences; only recognizing
# backticks let a '#'-prefixed line inside a ~~~-fenced block (e.g. a
# shell comment in a code sample) be misparsed as a real heading.
code_block_pattern = r"^(?:```|~~~)"
headers = []
in_code_block = False

View file

@ -77,7 +77,13 @@ class PdfParser:
try:
pix = pymupdf.Pixmap(image_bytes)
if pix.n > 4:
# n includes the alpha channel, so a plain RGBA pixmap also
# has n==4 — subtract alpha before comparing. Without this,
# a CMYK image with no alpha (n==4, same as RGBA) skips the
# RGB conversion, and pix.save() as .png then raises
# "unsupported colorspace for 'png'", silently dropping the
# image via the bare except below.
if pix.n - pix.alpha >= 4:
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
filename = f"p{page_num}_img{img_idx}.png"
save_path = images_path / filename

View file

@ -1,27 +1,24 @@
import json
import PyPDF2
try:
from .index.utils import get_number_of_pages, remove_fields, get_md_page_content
from .index.utils import (
get_number_of_pages, remove_fields, get_md_page_content,
parse_pages, get_pdf_page_content,
)
except ImportError:
from index.utils import get_number_of_pages, remove_fields, get_md_page_content
from index.utils import (
get_number_of_pages, remove_fields, get_md_page_content,
parse_pages, get_pdf_page_content,
)
# ── Helpers ──────────────────────────────────────────────────────────────────
def _parse_pages(pages: str) -> list[int]:
"""Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints."""
result = []
for part in pages.split(','):
part = part.strip()
if '-' in part:
start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip())
if start > end:
raise ValueError(f"Invalid range '{part}': start must be <= end")
result.extend(range(start, end + 1))
else:
result.append(int(part))
return sorted(set(result))
"""Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.
Delegates to the canonical implementation so the two never drift again
this one used to lack the p>=1 filter and the 1000-page DoS cap."""
return parse_pages(pages)
def _count_pages(doc_info: dict) -> int:
@ -34,7 +31,8 @@ def _count_pages(doc_info: dict) -> int:
def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
"""Extract text for specific PDF pages (1-indexed). Prefer cached pages, fallback to PDF."""
"""Extract text for specific PDF pages (1-indexed). Prefer cached pages,
else delegate the file-read fallback to the canonical implementation."""
cached_pages = doc_info.get('pages')
if cached_pages:
page_map = {p['page']: p['content'] for p in cached_pages}
@ -42,15 +40,7 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
{'page': p, 'content': page_map[p]}
for p in page_nums if p in page_map
]
path = doc_info['path']
with open(path, 'rb') as f:
pdf_reader = PyPDF2.PdfReader(f)
total = len(pdf_reader.pages)
valid_pages = [p for p in page_nums if 1 <= p <= total]
return [
{'page': p, 'content': pdf_reader.pages[p - 1].extract_text() or ''}
for p in valid_pages
]
return get_pdf_page_content(doc_info['path'], page_nums)
def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:

View file

@ -1,8 +1,22 @@
import json
import re
import sqlite3
import threading
from pathlib import Path
from ..errors import CollectionAlreadyExistsError, PageIndexError
# Mirrors LocalBackend's own collection-name rule. SQLiteStorage enforces this
# itself (not just relying on LocalBackend's pre-check or the schema's CHECK
# constraint below) because it's a public StorageEngine that can be used
# directly, bypassing LocalBackend entirely.
_COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
def _validate_collection_name(name: str) -> None:
if not _COLLECTION_NAME_RE.match(name):
raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
class SQLiteStorage:
def __init__(self, db_path: str):
@ -48,7 +62,17 @@ class SQLiteStorage:
conn.execute("PRAGMA user_version = 1")
conn.executescript("""
CREATE TABLE IF NOT EXISTS collections (
name TEXT PRIMARY KEY CHECK(length(name) <= 128 AND name GLOB '[a-zA-Z0-9_-]*'),
-- GLOB '*' is "any characters", not a regex quantifier over the
-- preceding class '[a-zA-Z0-9_-]*' alone only constrains the
-- FIRST character. The second GLOB (NOT ... '*[^...]*') checks
-- every remaining character too, so this is real defense-in-depth
-- for direct SQLiteStorage use (bypassing _validate_collection_name
-- above), not just a first-character gate.
name TEXT PRIMARY KEY CHECK(
length(name) BETWEEN 1 AND 128
AND name GLOB '[a-zA-Z0-9_-]*'
AND name NOT GLOB '*[^a-zA-Z0-9_-]*'
),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS documents (
@ -70,12 +94,17 @@ class SQLiteStorage:
conn.commit()
def create_collection(self, name: str) -> None:
_validate_collection_name(name)
with self._write_lock:
conn = self._get_conn()
conn.execute("INSERT INTO collections (name) VALUES (?)", (name,))
try:
conn.execute("INSERT INTO collections (name) VALUES (?)", (name,))
except sqlite3.IntegrityError as e:
raise CollectionAlreadyExistsError(f"Collection '{name}' already exists") from e
conn.commit()
def get_or_create_collection(self, name: str) -> None:
_validate_collection_name(name)
with self._write_lock:
conn = self._get_conn()
conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,))

View file

@ -14,13 +14,23 @@ class DocumentInfo(TypedDict):
doc_type: str
class DocumentDetail(DocumentInfo, total=False):
class _DocumentDetailRequired(DocumentInfo):
"""``structure`` is always present — split into its own (default
total=True) base so the total=False below only applies to the genuinely
optional, backend-specific fields below. A single
``class DocumentDetail(DocumentInfo, total=False): structure: ...`` would
incorrectly mark ``structure`` optional too, since total=False applies to
the whole class body, not just the fields declared after it.
"""
structure: list[dict[str, Any]]
class DocumentDetail(_DocumentDetailRequired, total=False):
"""A document with its tree, as returned by ``get_document()``.
``structure`` is always present; ``file_path`` is local-only and
``status`` is cloud-only, hence total=False.
``status`` is cloud-only, hence total=False for those two only.
"""
structure: list[dict[str, Any]]
file_path: str # local backend only
status: str # cloud backend only