Coverage for src/lilbee/data/extract/xberg.py: 100%

1"""Bridge to xberg's async-only ``extract`` for lilbee's call sites. 

2 

3xberg exposes one ``extract(input, config)`` coroutine; lilbee extracts a single 

4in-memory document at a time, from both async and sync callers. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from concurrent.futures import ThreadPoolExecutor 

11from dataclasses import dataclass 

12from typing import TYPE_CHECKING 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Coroutine 

16 

17 from xberg import ( 

18 ExtractedDocument, 

19 ExtractInput, 

20 ExtractionConfig, 

21 ExtractionResult, 

22 OcrConfig, 

23 ) 

24 

25 

26@dataclass(frozen=True) 

27class BatchItem: 

28 """One input for :func:`aextract_batch`, with its per-file OCR override.""" 

29 

30 data: bytes 

31 mime: str | None 

32 filename: str | None 

33 ocr: OcrConfig | None 

34 

35 

36def _input(data: bytes, mime_type: str | None, filename: str | None) -> ExtractInput: 

37 from xberg import ExtractInput, ExtractInputKind 

38 

39 return ExtractInput( 

40 kind=ExtractInputKind.BYTES, bytes=data, mime_type=mime_type, filename=filename 

41 ) 

42 

43 

44def _first(result: ExtractionResult) -> ExtractedDocument: 

45 """Return the single extracted document, or raise on an extraction error.""" 

46 if result.results: 

47 return result.results[0] 

48 if result.errors: 

49 raise RuntimeError(str(result.errors[0])) 

50 raise RuntimeError("xberg extraction returned no document") 

51 

52 

53async def aextract_document( 

54 data: bytes, 

55 mime_type: str | None = None, 

56 *, 

57 filename: str | None = None, 

58 config: ExtractionConfig, 

59) -> ExtractedDocument: 

60 """Extract one in-memory document. For callers already on the event loop.""" 

61 from xberg import extract 

62 

63 return _first(await extract(_input(data, mime_type, filename), config)) 

64 

65 

66async def aextract_batch( 

67 items: list[BatchItem], config: ExtractionConfig 

68) -> list[ExtractedDocument | Exception]: 

69 """Extract many inputs in one call, returning one document-or-error per input. 

70 

71 Each item's OCR config overrides the batch default for that file. xberg compacts 

72 ``results`` to successes in input order and reports failures in ``errors`` by 

73 input index; this remaps them back to one slot per input. 

74 """ 

75 from xberg import ExtractInput, ExtractInputKind, FileExtractionConfig, extract_batch 

76 

77 inputs = [ 

78 ExtractInput( 

79 kind=ExtractInputKind.BYTES, 

80 bytes=item.data, 

81 mime_type=item.mime, 

82 filename=item.filename, 

83 config=FileExtractionConfig(ocr=item.ocr) if item.ocr is not None else None, 

84 ) 

85 for item in items 

86 ] 

87 result = await extract_batch(inputs, config) 

88 failed: dict[int, Exception] = {e.index: RuntimeError(e.message) for e in result.errors} 

89 success_indices = [i for i in range(len(items)) if i not in failed] 

90 by_index: dict[int, ExtractedDocument | Exception] = dict( 

91 zip(success_indices, result.results, strict=True) 

92 ) 

93 by_index.update(failed) 

94 return [by_index[i] for i in range(len(items))] 

95 

96 

97def extract_document( 

98 data: bytes, 

99 mime_type: str | None = None, 

100 *, 

101 filename: str | None = None, 

102 config: ExtractionConfig, 

103) -> ExtractedDocument: 

104 """Extract one in-memory document from synchronous code. 

105 

106 Uses ``asyncio.run``; if a loop is already running on this thread, drives the 

107 coroutine on a fresh worker thread so it never re-enters that loop. 

108 """ 

109 return _run(aextract_document(data, mime_type, filename=filename, config=config)) 

110 

111 

112def _run(coro: Coroutine[None, None, ExtractedDocument]) -> ExtractedDocument: 

113 try: 

114 asyncio.get_running_loop() 

115 except RuntimeError: 

116 return asyncio.run(coro) 

117 with ThreadPoolExecutor(max_workers=1) as pool: 

118 return pool.submit(asyncio.run, coro).result()