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 

47 The error item carries the reason in ``message``; it has no ``__str__``, so 

48 formatting the item itself would hand the caller an object repr instead of 

49 the timeout or unsupported-format it is reporting. 

50 """ 

51 if result.results: 

52 return result.results[0] 

53 if result.errors: 

54 raise RuntimeError(result.errors[0].message) 

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

56 

57 

58async def aextract_document( 

59 data: bytes, 

60 mime_type: str | None = None, 

61 *, 

62 filename: str | None = None, 

63 config: ExtractionConfig, 

64) -> ExtractedDocument: 

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

66 from xberg import extract 

67 

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

69 

70 

71async def aextract_batch( 

72 items: list[BatchItem], config: ExtractionConfig 

73) -> list[ExtractedDocument | Exception]: 

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

75 

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

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

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

79 """ 

80 from xberg import ExtractInput, ExtractInputKind, FileExtractionConfig, extract_batch 

81 

82 inputs = [ 

83 ExtractInput( 

84 kind=ExtractInputKind.BYTES, 

85 bytes=item.data, 

86 mime_type=item.mime, 

87 filename=item.filename, 

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

89 ) 

90 for item in items 

91 ] 

92 result = await extract_batch(inputs, config) 

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

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

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

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

97 ) 

98 by_index.update(failed) 

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

100 

101 

102def extract_document( 

103 data: bytes, 

104 mime_type: str | None = None, 

105 *, 

106 filename: str | None = None, 

107 config: ExtractionConfig, 

108) -> ExtractedDocument: 

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

110 

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

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

113 """ 

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

115 

116 

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

118 try: 

119 asyncio.get_running_loop() 

120 except RuntimeError: 

121 return asyncio.run(coro) 

122 with ThreadPoolExecutor(max_workers=1) as pool: 

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