Coverage for src / quber / core / extractors / dual / models.py: 100%

1"""Result types for the two-flow extractor. 

2 

3`VisualTable` is one table as identified by looking at the page (the 

4authoritative count and identity). `CamelotTable` is one cell grid pulled 

5by Camelot, carrying its classifier verdict as metadata only — the verdict 

6never drops a candidate here, because identity is the visual flow's job. 

7`DualResult` holds both lists for one document, unmatched. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import List, Optional, Tuple 

13 

14from pydantic import BaseModel, Field 

15 

16from quber.agents.classifier import ClassifierResult 

17from quber.core.extractors.camelot.acquire import CamelotCandidate 

18 

19 

20class VisualTable(BaseModel): 

21 """A table the visual flow found on a page, with its full-anatomy box.""" 

22 

23 page: int = Field(ge=1, description="1-indexed page the table sits on") 

24 ordinal: int = Field(ge=1, description="1-based reading-order position within the page") 

25 title: str = Field(default="", description="Visible caption, else a one-line column summary") 

26 region: Tuple[float, float, float, float] = Field( 

27 description="Boundary box, normalized 0..1 with the page top-left as origin (x1, y1, x2, y2)" 

28 ) 

29 tightened: bool = Field( 

30 description="True if the box was snapped to the text layer; False if the coarse grid cell was kept" 

31 ) 

32 

33 

34class CamelotTable(BaseModel): 

35 """A Camelot cell grid plus its is-table verdict (metadata, not a gate).""" 

36 

37 candidate: CamelotCandidate = Field(description="Raw Camelot output: cells, markdown, bbox, accuracy") 

38 classification: Optional[ClassifierResult] = Field( 

39 default=None, 

40 description="Classifier's is_table verdict, kept for later cross-reference; it does not drop the candidate", 

41 ) 

42 

43 

44class DualResult(BaseModel): 

45 """Both flows' output for one document, side by side and unmatched.""" 

46 

47 source: str = Field(description="Path of the source PDF") 

48 visual: List[VisualTable] = Field( 

49 default_factory=list, description="Tables identified by the visual flow" 

50 ) 

51 camelot: List[CamelotTable] = Field(default_factory=list, description="Cell grids extracted by Camelot") 

52 

53 def visual_count_by_page(self) -> dict[int, int]: 

54 counts: dict[int, int] = {} 

55 for t in self.visual: 

56 counts[t.page] = counts.get(t.page, 0) + 1 

57 return counts 

58 

59 def camelot_count_by_page(self) -> dict[int, int]: 

60 counts: dict[int, int] = {} 

61 for t in self.camelot: 

62 counts[t.candidate.page] = counts.get(t.candidate.page, 0) + 1 

63 return counts