chore(lint): fix ruff violations on Neo4j files

Pre-existing ruff violations surfaced by CI on PR #188:

- TC001: move runtime `Neo4jDriver` imports into `TYPE_CHECKING`
  blocks (queries.py, chunk_writer.py, schema.py, tree_reader.py,
  tree_writer.py)
- SIM117: combine nested `async with` in chunk_writer.write_chunks
  and tree_writer.write_document
- SIM105: replace try/except/pass with `contextlib.suppress` in
  tree_writer._element_props
- F401: remove unused `Neo4jDriver` import in main._init_neo4j
- RUF100: remove unused `# noqa: E402` in tests/neo4j/conftest.py
- I001: sort imports in tests/neo4j/test_chunk_writer.py and
  test_tree_writer.py

Zero behaviour change. `ruff check .` now passes cleanly.
This commit is contained in:
Pier-Jean Malandrino 2026-04-20 10:22:33 +02:00
parent 03e5dfac17
commit bd73c4bbfd
9 changed files with 209 additions and 193 deletions

View file

@ -16,9 +16,10 @@ from __future__ import annotations
import json import json
import logging import logging
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any from typing import TYPE_CHECKING, Any
from infra.neo4j.driver import Neo4jDriver if TYPE_CHECKING:
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -63,23 +64,23 @@ async def write_chunks(
if ref: if ref:
derived_rows.append({"chunk_id": cid, "doc_id": doc_id, "self_ref": ref}) derived_rows.append({"chunk_id": cid, "doc_id": doc_id, "self_ref": ref})
async with neo.driver.session(database=neo.database) as session: async with (
async with await session.begin_transaction() as tx: neo.driver.session(database=neo.database) as session,
# Replace existing chunks. await session.begin_transaction() as tx,
await tx.run( ):
""" # Replace existing chunks.
await tx.run(
"""
MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk) MATCH (d:Document {id: $doc_id})-[:HAS_CHUNK]->(c:Chunk)
DETACH DELETE c DETACH DELETE c
""", """,
doc_id=doc_id, doc_id=doc_id,
) )
await tx.run( await tx.run("MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id)
"MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id
)
if chunk_rows: if chunk_rows:
await tx.run( await tx.run(
""" """
MATCH (d:Document {id: $doc_id}) MATCH (d:Document {id: $doc_id})
UNWIND $rows AS r UNWIND $rows AS r
CREATE (c:Chunk { CREATE (c:Chunk {
@ -92,33 +93,33 @@ async def write_chunks(
}) })
MERGE (d)-[:HAS_CHUNK]->(c) MERGE (d)-[:HAS_CHUNK]->(c)
""", """,
doc_id=doc_id, doc_id=doc_id,
rows=chunk_rows, rows=chunk_rows,
) )
if derived_rows: if derived_rows:
await tx.run( await tx.run(
""" """
UNWIND $rows AS r UNWIND $rows AS r
MATCH (c:Chunk {id: r.chunk_id}) MATCH (c:Chunk {id: r.chunk_id})
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref}) MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
MERGE (c)-[:DERIVED_FROM]->(e) MERGE (c)-[:DERIVED_FROM]->(e)
""", """,
rows=derived_rows, rows=derived_rows,
) )
# Flag the Document with the new stage. # Flag the Document with the new stage.
await tx.run( await tx.run(
""" """
MATCH (d:Document {id: $doc_id}) MATCH (d:Document {id: $doc_id})
SET d.stages_applied = [s IN coalesce(d.stages_applied, []) WHERE s <> 'chunks'] SET d.stages_applied = [s IN coalesce(d.stages_applied, []) WHERE s <> 'chunks']
+ ['chunks'], + ['chunks'],
d.last_chunk_write = datetime() d.last_chunk_write = datetime()
""", """,
doc_id=doc_id, doc_id=doc_id,
) )
await tx.commit() await tx.commit()
logger.info( logger.info(
"Neo4j: wrote %d chunks (%d DERIVED_FROM) for doc %s", "Neo4j: wrote %d chunks (%d DERIVED_FROM) for doc %s",

View file

@ -3,9 +3,10 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any from typing import TYPE_CHECKING, Any
from infra.neo4j.driver import Neo4jDriver if TYPE_CHECKING:
from infra.neo4j.driver import Neo4jDriver
@dataclass @dataclass

View file

@ -7,21 +7,21 @@ this multiple times is safe — it's the contract integration tests rely on.
from __future__ import annotations from __future__ import annotations
import logging import logging
from typing import TYPE_CHECKING
from infra.neo4j.driver import Neo4jDriver if TYPE_CHECKING:
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
CONSTRAINTS: tuple[str, ...] = ( CONSTRAINTS: tuple[str, ...] = (
"CREATE CONSTRAINT document_id IF NOT EXISTS " "CREATE CONSTRAINT document_id IF NOT EXISTS FOR (d:Document) REQUIRE d.id IS UNIQUE",
"FOR (d:Document) REQUIRE d.id IS UNIQUE",
"CREATE CONSTRAINT element_composite IF NOT EXISTS " "CREATE CONSTRAINT element_composite IF NOT EXISTS "
"FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE", "FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE",
"CREATE CONSTRAINT page_composite IF NOT EXISTS " "CREATE CONSTRAINT page_composite IF NOT EXISTS "
"FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE", "FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE",
"CREATE CONSTRAINT chunk_id IF NOT EXISTS " "CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
"FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
) )
INDEXES: tuple[str, ...] = ( INDEXES: tuple[str, ...] = (
@ -30,8 +30,7 @@ INDEXES: tuple[str, ...] = (
) )
FULLTEXT_INDEXES: tuple[str, ...] = ( FULLTEXT_INDEXES: tuple[str, ...] = (
"CREATE FULLTEXT INDEX element_text IF NOT EXISTS " "CREATE FULLTEXT INDEX element_text IF NOT EXISTS FOR (e:Element) ON EACH [e.text]",
"FOR (e:Element) ON EACH [e.text]",
) )

View file

@ -10,8 +10,10 @@ nodes directly.
from __future__ import annotations from __future__ import annotations
import logging import logging
from typing import TYPE_CHECKING
from infra.neo4j.driver import Neo4jDriver if TYPE_CHECKING:
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -57,10 +59,6 @@ async def delete_document(neo: Neo4jDriver, doc_id: str) -> int:
) )
record = await result.single() record = await result.single()
# Also clean up orphan elements and pages tagged with this doc_id. # Also clean up orphan elements and pages tagged with this doc_id.
await session.run( await session.run("MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id)
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id await session.run("MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id)
)
await session.run(
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
)
return int(record["removed"]) if record else 0 return int(record["removed"]) if record else 0

View file

@ -9,13 +9,15 @@ graph nodes is deferred to v0.6 (see docs/design/neo4j-integration.md §2).
from __future__ import annotations from __future__ import annotations
import contextlib
import json import json
import logging import logging
from dataclasses import dataclass from dataclasses import dataclass
from datetime import datetime, timezone from datetime import UTC, datetime
from typing import Any from typing import TYPE_CHECKING, Any
from infra.neo4j.driver import Neo4jDriver if TYPE_CHECKING:
from infra.neo4j.driver import Neo4jDriver
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -96,10 +98,8 @@ def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]:
props["caption"] = item.get("caption") props["caption"] = item.get("caption")
if item.get("data") and isinstance(item["data"], dict): if item.get("data") and isinstance(item["data"], dict):
# Tables carry cell layout under data; stringify to keep the schema flat. # Tables carry cell layout under data; stringify to keep the schema flat.
try: with contextlib.suppress(TypeError, ValueError):
props["cells_json"] = json.dumps(item["data"]) props["cells_json"] = json.dumps(item["data"])
except (TypeError, ValueError):
pass
return props return props
@ -145,7 +145,7 @@ async def write_document(
Fails fast (exception propagates) if Neo4j is unavailable per design §8.5. Fails fast (exception propagates) if Neo4j is unavailable per design §8.5.
""" """
doc_data = json.loads(document_json) doc_data = json.loads(document_json)
ingested_at = datetime.now(tz=timezone.utc).isoformat() ingested_at = datetime.now(tz=UTC).isoformat()
elements: list[dict[str, Any]] = [] elements: list[dict[str, Any]] = []
for _, item in _iter_items(doc_data): for _, item in _iter_items(doc_data):
@ -179,26 +179,24 @@ async def write_document(
reading_order = _dfs_order(doc_data) reading_order = _dfs_order(doc_data)
async with neo.driver.session(database=neo.database) as session: async with (
async with await session.begin_transaction() as tx: neo.driver.session(database=neo.database) as session,
# 1. Wipe existing graph for this doc_id (replace strategy). await session.begin_transaction() as tx,
await tx.run( ):
"MATCH (d:Document {id: $doc_id}) " # 1. Wipe existing graph for this doc_id (replace strategy).
"OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n) " await tx.run(
"DETACH DELETE d, n", "MATCH (d:Document {id: $doc_id}) "
doc_id=doc_id, "OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK*0..]->(n) "
) "DETACH DELETE d, n",
# Also wipe orphan elements/chunks that may still reference this doc. doc_id=doc_id,
await tx.run( )
"MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id # Also wipe orphan elements/chunks that may still reference this doc.
) await tx.run("MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id)
await tx.run( await tx.run("MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id)
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
)
# 2. Document node (carries the verbatim JSON for TreeReader). # 2. Document node (carries the verbatim JSON for TreeReader).
await tx.run( await tx.run(
""" """
CREATE (d:Document { CREATE (d:Document {
id: $doc_id, id: $doc_id,
title: $title, title: $title,
@ -211,32 +209,32 @@ async def write_document(
document_json: $document_json document_json: $document_json
}) })
""", """,
doc_id=doc_id, doc_id=doc_id,
title=filename, title=filename,
source_uri=source_uri or "", source_uri=source_uri or "",
ingested_at=ingested_at, ingested_at=ingested_at,
docling_version=docling_version or "", docling_version=docling_version or "",
tenant_id=tenant_id, tenant_id=tenant_id,
document_json=document_json, document_json=document_json,
)
# 3. Page nodes.
if pages:
await tx.run(
"UNWIND $pages AS p "
"CREATE (:Page {doc_id: p.doc_id, page_no: p.page_no, "
"width: p.width, height: p.height})",
pages=pages,
) )
# 3. Page nodes. # 4. Element nodes — use dynamic :Element:<specific> labels via APOC-free trick.
if pages: # We split by specific label so the CREATE statement is static (no APOC).
await tx.run( by_specific: dict[str, list[dict[str, Any]]] = {}
"UNWIND $pages AS p " for e in elements:
"CREATE (:Page {doc_id: p.doc_id, page_no: p.page_no, " by_specific.setdefault(e["specific_label"], []).append(e)
"width: p.width, height: p.height})", for specific, batch in by_specific.items():
pages=pages, await tx.run(
) f"""
# 4. Element nodes — use dynamic :Element:<specific> labels via APOC-free trick.
# We split by specific label so the CREATE statement is static (no APOC).
by_specific: dict[str, list[dict[str, Any]]] = {}
for e in elements:
by_specific.setdefault(e["specific_label"], []).append(e)
for specific, batch in by_specific.items():
await tx.run(
f"""
UNWIND $batch AS e UNWIND $batch AS e
CREATE (n:Element:{specific} {{ CREATE (n:Element:{specific} {{
doc_id: e.doc_id, doc_id: e.doc_id,
@ -250,83 +248,83 @@ async def write_document(
cells_json: e.cells_json cells_json: e.cells_json
}}) }})
""", """,
batch=batch, batch=batch,
) )
# 5. PARENT_OF relations (tree structure). Order tracked inline. # 5. PARENT_OF relations (tree structure). Order tracked inline.
parent_rows = [ parent_rows = [
{ {
"doc_id": doc_id, "doc_id": doc_id,
"parent_ref": e["parent_ref"], "parent_ref": e["parent_ref"],
"child_ref": e["self_ref"], "child_ref": e["self_ref"],
"order": idx, "order": idx,
} }
for idx, e in enumerate(elements) for idx, e in enumerate(elements)
if e["parent_ref"] and e["parent_ref"] != "#/body" if e["parent_ref"] and e["parent_ref"] != "#/body"
] ]
if parent_rows: if parent_rows:
await tx.run( await tx.run(
""" """
UNWIND $rows AS r UNWIND $rows AS r
MATCH (p:Element {doc_id: r.doc_id, self_ref: r.parent_ref}) MATCH (p:Element {doc_id: r.doc_id, self_ref: r.parent_ref})
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref}) MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
MERGE (p)-[rel:PARENT_OF]->(c) MERGE (p)-[rel:PARENT_OF]->(c)
SET rel.order = r.order SET rel.order = r.order
""", """,
rows=parent_rows, rows=parent_rows,
) )
# 6. HAS_ROOT for top-level children of the document body. # 6. HAS_ROOT for top-level children of the document body.
root_rows = [ root_rows = [
{"doc_id": doc_id, "child_ref": e["self_ref"]} {"doc_id": doc_id, "child_ref": e["self_ref"]}
for e in elements for e in elements
if e["parent_ref"] == "#/body" if e["parent_ref"] == "#/body"
] ]
if root_rows: if root_rows:
await tx.run( await tx.run(
""" """
UNWIND $rows AS r UNWIND $rows AS r
MATCH (d:Document {id: r.doc_id}) MATCH (d:Document {id: r.doc_id})
MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref}) MATCH (c:Element {doc_id: r.doc_id, self_ref: r.child_ref})
MERGE (d)-[:HAS_ROOT]->(c) MERGE (d)-[:HAS_ROOT]->(c)
""", """,
rows=root_rows, rows=root_rows,
) )
# 7. ON_PAGE from first provenance. # 7. ON_PAGE from first provenance.
on_page_rows = [ on_page_rows = [
{"doc_id": doc_id, "self_ref": e["self_ref"], "page_no": e["prov_page"]} {"doc_id": doc_id, "self_ref": e["self_ref"], "page_no": e["prov_page"]}
for e in elements for e in elements
if e["prov_page"] is not None if e["prov_page"] is not None
] ]
if on_page_rows: if on_page_rows:
await tx.run( await tx.run(
""" """
UNWIND $rows AS r UNWIND $rows AS r
MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref}) MATCH (e:Element {doc_id: r.doc_id, self_ref: r.self_ref})
MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no}) MATCH (p:Page {doc_id: r.doc_id, page_no: r.page_no})
MERGE (e)-[:ON_PAGE]->(p) MERGE (e)-[:ON_PAGE]->(p)
""", """,
rows=on_page_rows, rows=on_page_rows,
) )
# 8. NEXT chain in DFS pre-order. # 8. NEXT chain in DFS pre-order.
if len(reading_order) > 1: if len(reading_order) > 1:
pairs = [ pairs = [
{"doc_id": doc_id, "a": reading_order[i], "b": reading_order[i + 1]} {"doc_id": doc_id, "a": reading_order[i], "b": reading_order[i + 1]}
for i in range(len(reading_order) - 1) for i in range(len(reading_order) - 1)
] ]
await tx.run( await tx.run(
""" """
UNWIND $pairs AS p UNWIND $pairs AS p
MATCH (a:Element {doc_id: p.doc_id, self_ref: p.a}) MATCH (a:Element {doc_id: p.doc_id, self_ref: p.a})
MATCH (b:Element {doc_id: p.doc_id, self_ref: p.b}) MATCH (b:Element {doc_id: p.doc_id, self_ref: p.b})
MERGE (a)-[:NEXT]->(b) MERGE (a)-[:NEXT]->(b)
""", """,
pairs=pairs, pairs=pairs,
) )
await tx.commit() await tx.commit()
logger.info( logger.info(
"Neo4j: wrote doc %s (%d elements, %d pages)", "Neo4j: wrote doc %s (%d elements, %d pages)",

View file

@ -101,7 +101,7 @@ async def _init_neo4j():
logger.info("Neo4j disabled (NEO4J_URI not set)") logger.info("Neo4j disabled (NEO4J_URI not set)")
return None return None
from infra.neo4j import Neo4jDriver, bootstrap_schema, get_driver from infra.neo4j import bootstrap_schema, get_driver
try: try:
neo = await get_driver( neo = await get_driver(

View file

@ -15,7 +15,7 @@ import pytest
# (e.g. local dev without the dependency installed). # (e.g. local dev without the dependency installed).
pytest.importorskip("neo4j") pytest.importorskip("neo4j")
from infra.neo4j import close_driver, get_driver # noqa: E402 from infra.neo4j import close_driver, get_driver
def _cfg() -> tuple[str, str, str]: def _cfg() -> tuple[str, str, str]:

View file

@ -12,7 +12,6 @@ from infra.neo4j import fetch_graph, write_chunks, write_document
from infra.neo4j.schema import bootstrap_schema from infra.neo4j.schema import bootstrap_schema
from tests.neo4j.test_tree_writer import FIXTURE from tests.neo4j.test_tree_writer import FIXTURE
CHUNKS = [ CHUNKS = [
{ {
"text": "Introduction. First paragraph on page 1.", "text": "Introduction. First paragraph on page 1.",
@ -66,8 +65,7 @@ async def test_write_chunks_and_derived_from(neo4j_driver):
for idx, expected in [(0, 2), (1, 1)]: for idx, expected in [(0, 2), (1, 1)]:
cnt = await ( cnt = await (
await s.run( await s.run(
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) " "MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) RETURN count(e) AS n",
"RETURN count(e) AS n",
cid=f"doc-fixture::chunk::{idx}", cid=f"doc-fixture::chunk::{idx}",
) )
).single() ).single()

View file

@ -13,7 +13,6 @@ import json
from infra.neo4j import read_document_json, write_document from infra.neo4j import read_document_json, write_document
from infra.neo4j.schema import bootstrap_schema from infra.neo4j.schema import bootstrap_schema
FIXTURE = { FIXTURE = {
"name": "fixture.pdf", "name": "fixture.pdf",
"pages": { "pages": {
@ -89,41 +88,59 @@ async def test_write_creates_expected_structure(neo4j_driver):
assert result.pages_written == 2 assert result.pages_written == 2
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s: async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
assert await _count( assert (
s, await _count(
"MATCH (d:Document {id: $id}) RETURN count(d) AS n", s,
id="doc-fixture", "MATCH (d:Document {id: $id}) RETURN count(d) AS n",
) == 1 id="doc-fixture",
assert await _count( )
s, == 1
"MATCH (:Document {id: $id})-[:HAS_ROOT]->(e:Element) RETURN count(e) AS n", )
id="doc-fixture", assert (
) == 4 await _count(
assert await _count( s,
s, "MATCH (:Document {id: $id})-[:HAS_ROOT]->(e:Element) RETURN count(e) AS n",
"MATCH (e:Element:SectionHeader {doc_id: $id, self_ref: '#/texts/0'}) " id="doc-fixture",
"RETURN count(e) AS n", )
id="doc-fixture", == 4
) == 1 )
assert await _count( assert (
s, await _count(
"MATCH (e:Element:Table {doc_id: $id}) RETURN count(e) AS n", s,
id="doc-fixture", "MATCH (e:Element:SectionHeader {doc_id: $id, self_ref: '#/texts/0'}) "
) == 1 "RETURN count(e) AS n",
id="doc-fixture",
)
== 1
)
assert (
await _count(
s,
"MATCH (e:Element:Table {doc_id: $id}) RETURN count(e) AS n",
id="doc-fixture",
)
== 1
)
# Reading-order chain: 3 NEXT edges for 4 elements. # Reading-order chain: 3 NEXT edges for 4 elements.
assert await _count( assert (
s, await _count(
"MATCH (a:Element {doc_id: $id})-[:NEXT]->(b:Element {doc_id: $id}) " s,
"RETURN count(*) AS n", "MATCH (a:Element {doc_id: $id})-[:NEXT]->(b:Element {doc_id: $id}) "
id="doc-fixture", "RETURN count(*) AS n",
) == 3 id="doc-fixture",
)
== 3
)
# ON_PAGE: one per element with prov. # ON_PAGE: one per element with prov.
assert await _count( assert (
s, await _count(
"MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) " s,
"RETURN count(*) AS n", "MATCH (:Element {doc_id: $id})-[:ON_PAGE]->(:Page {doc_id: $id}) "
id="doc-fixture", "RETURN count(*) AS n",
) == 4 id="doc-fixture",
)
== 4
)
async def test_rewrite_is_idempotent_replace(neo4j_driver): async def test_rewrite_is_idempotent_replace(neo4j_driver):
@ -145,14 +162,18 @@ async def test_rewrite_is_idempotent_replace(neo4j_driver):
) )
async with neo4j_driver.driver.session(database=neo4j_driver.database) as s: async with neo4j_driver.driver.session(database=neo4j_driver.database) as s:
assert await _count( assert (
s, "MATCH (d:Document {id: $id}) RETURN count(d) AS n", id="doc-fixture" await _count(s, "MATCH (d:Document {id: $id}) RETURN count(d) AS n", id="doc-fixture")
) == 1 == 1
assert await _count( )
s, assert (
"MATCH (e:Element {doc_id: $id}) RETURN count(e) AS n", await _count(
id="doc-fixture", s,
) == 4 "MATCH (e:Element {doc_id: $id}) RETURN count(e) AS n",
id="doc-fixture",
)
== 4
)
async def test_reader_returns_verbatim_json(neo4j_driver): async def test_reader_returns_verbatim_json(neo4j_driver):