Coverage for src/local_deep_research/document_loaders/toml_loader.py: 100%

1"""Load TOML with Python's bundled parser, within the shared TOML limits.""" 

2 

3import json 

4from pathlib import Path 

5from typing import Iterator 

6 

7from langchain_core.document_loaders import BaseLoader 

8from langchain_core.documents import Document 

9 

10from ..security import toml_limits 

11 

12 

13def _read_limited_text(path: Path) -> str: 

14 """Return *path*'s text without reading past the TOML size limit.""" 

15 with path.open("rb") as stream: 

16 content = stream.read(toml_limits.MAX_TOML_BYTES + 1) 

17 # Checked before decoding, so a read cut off inside a multi-byte 

18 # character is reported as too large rather than as invalid UTF-8. 

19 toml_limits.check_toml_size(len(content)) 

20 return content.decode() 

21 

22 

23class TOMLLoader(BaseLoader): 

24 """Load a TOML file or directory without requiring the tomli backport. 

25 

26 Each file is refused with TOMLLimitError, before tomllib parses it, when 

27 it is larger than MAX_TOML_BYTES or has a key or table header of more 

28 than MAX_TOML_KEY_PARTS dotted parts (see security.toml_limits). 

29 """ 

30 

31 def __init__(self, source: str | Path): 

32 # nosemgrep: semgrep.rules.path-traversal-risk, reason: load_from_bytes passes the temp file it wrote; no route passes a request path 

33 self.source = Path(source) 

34 

35 def lazy_load(self) -> Iterator[Document]: 

36 if self.source.is_file() and self.source.suffix.lower() == ".toml": 

37 files = [self.source] 

38 elif self.source.is_dir(): 

39 files = sorted(self.source.rglob("*.toml")) 

40 else: 

41 raise ValueError("Invalid source path or file type") 

42 

43 for file_path in files: 

44 data = toml_limits.loads_toml_within_limits( 

45 _read_limited_text(file_path) 

46 ) 

47 yield Document( 

48 page_content=json.dumps(data, ensure_ascii=False, default=str), 

49 metadata={"source": str(file_path)}, 

50 )