danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

chore(lint): fix ruff violations on Neo4j files

Browse source 
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
Show stats Download patch file Download diff file Expand all files Collapse all files

15
document-parser/infra/neo4j/chunk_writer.py
View file

@ -16,9 +16,10 @@ from __future__ import annotations
import json
import logging
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__)
@ -63,8 +64,10 @@ async def write_chunks(
if 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 await session.begin_transaction() as tx:
async with (
neo.driver.session(database=neo.database) as session,
await session.begin_transaction() as tx,
):
# Replace existing chunks.
await tx.run(
"""
@ -73,9 +76,7 @@ async def write_chunks(
""",
doc_id=doc_id,
)
await tx.run(
"MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id
)
await tx.run("MATCH (c:Chunk {doc_id: $doc_id}) DETACH DELETE c", doc_id=doc_id)
if chunk_rows:
await tx.run(

5
document-parser/infra/neo4j/queries.py
View file

@ -3,9 +3,10 @@
from __future__ import annotations
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

13
document-parser/infra/neo4j/schema.py
View file

@ -7,21 +7,21 @@ this multiple times is safe — it's the contract integration tests rely on.
from __future__ import annotations
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__)
CONSTRAINTS: tuple[str, ...] = (
"CREATE CONSTRAINT document_id IF NOT EXISTS "
"FOR (d:Document) REQUIRE d.id IS UNIQUE",
"CREATE CONSTRAINT document_id IF NOT EXISTS FOR (d:Document) REQUIRE d.id IS UNIQUE",
"CREATE CONSTRAINT element_composite IF NOT EXISTS "
"FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE",
"CREATE CONSTRAINT page_composite IF NOT EXISTS "
"FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE",
"CREATE CONSTRAINT chunk_id IF NOT EXISTS "
"FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
"CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE",
)
INDEXES: tuple[str, ...] = (
@ -30,8 +30,7 @@ INDEXES: tuple[str, ...] = (
)
FULLTEXT_INDEXES: tuple[str, ...] = (
"CREATE FULLTEXT INDEX element_text IF NOT EXISTS "
"FOR (e:Element) ON EACH [e.text]",
"CREATE FULLTEXT INDEX element_text IF NOT EXISTS FOR (e:Element) ON EACH [e.text]",
)

12
document-parser/infra/neo4j/tree_reader.py
View file

@ -10,8 +10,10 @@ nodes directly.
from __future__ import annotations
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__)
@ -57,10 +59,6 @@ async def delete_document(neo: Neo4jDriver, doc_id: str) -> int:
)
record = await result.single()
# Also clean up orphan elements and pages tagged with this doc_id.
await session.run(
"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 (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)
return int(record["removed"]) if record else 0

28
document-parser/infra/neo4j/tree_writer.py
View file

@ -9,13 +9,15 @@ graph nodes is deferred to v0.6 (see docs/design/neo4j-integration.md §2).
from __future__ import annotations
import contextlib
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from datetime import UTC, datetime
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__)
@ -96,10 +98,8 @@ def _element_props(item: dict[str, Any], doc_id: str) -> dict[str, Any]:
props["caption"] = item.get("caption")
if item.get("data") and isinstance(item["data"], dict):
# 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"])
except (TypeError, ValueError):
pass
return props
@ -145,7 +145,7 @@ async def write_document(
Fails fast (exception propagates) if Neo4j is unavailable — per design §8.5.
"""
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]] = []
for _, item in _iter_items(doc_data):
@ -179,8 +179,10 @@ async def write_document(
reading_order = _dfs_order(doc_data)
async with neo.driver.session(database=neo.database) as session:
async with await session.begin_transaction() as tx:
async with (
neo.driver.session(database=neo.database) as session,
await session.begin_transaction() as tx,
):
# 1. Wipe existing graph for this doc_id (replace strategy).
await tx.run(
"MATCH (d:Document {id: $doc_id}) "
@ -189,12 +191,8 @@ async def write_document(
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(
"MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id
)
await tx.run("MATCH (e:Element {doc_id: $doc_id}) DETACH DELETE e", doc_id=doc_id)
await tx.run("MATCH (p:Page {doc_id: $doc_id}) DETACH DELETE p", doc_id=doc_id)
# 2. Document node (carries the verbatim JSON for TreeReader).
await tx.run(

2
document-parser/main.py
View file

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

2
document-parser/tests/neo4j/conftest.py
View file

@ -15,7 +15,7 @@ import pytest
# (e.g. local dev without the dependency installed).
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]:

4
document-parser/tests/neo4j/test_chunk_writer.py
View file

@ -12,7 +12,6 @@ from infra.neo4j import fetch_graph, write_chunks, write_document
from infra.neo4j.schema import bootstrap_schema
from tests.neo4j.test_tree_writer import FIXTURE
CHUNKS = [
{
"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)]:
cnt = await (
await s.run(
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) "
"RETURN count(e) AS n",
"MATCH (c:Chunk {id: $cid})-[:DERIVED_FROM]->(e:Element) RETURN count(e) AS n",
cid=f"doc-fixture::chunk::{idx}",
)
).single()

57
document-parser/tests/neo4j/test_tree_writer.py
View file

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