Coverage for src/local_deep_research/document_loaders/evernote_loader.py: 96%

1"""Read Evernote exports using the project's existing XML and HTML parsers.""" 

2 

3from pathlib import Path 

4from time import strptime 

5from typing import Any, Iterator 

6 

7# nosemgrep: python.lang.security.use-defused-xml.use-defused-xml, reason: TreeBuilder only assembles elements; defusedxml's parser reads the XML 

8from xml.etree.ElementTree import TreeBuilder 

9 

10from bs4 import BeautifulSoup 

11from defusedxml.common import DTDForbidden 

12from defusedxml.ElementTree import DefusedXMLParser, iterparse 

13from langchain_core.document_loaders import BaseLoader 

14from langchain_core.documents import Document 

15 

16 

17class _EvernoteXMLParser(DefusedXMLParser): 

18 """defusedxml's parser, refusing any DOCTYPE with an internal subset. 

19 

20 Evernote exports declare ``<!DOCTYPE en-export SYSTEM "...">``, and 

21 defusedxml's ``forbid_dtd`` refuses every DOCTYPE, that one included. 

22 This parser keeps ``forbid_dtd`` on but refuses only a DOCTYPE with an 

23 internal subset (``[...]``), where a file could declare entities or 

24 attribute defaults that expat adds to every matching element. Expat 

25 reports the subset when it reaches the ``[``, so it is refused before 

26 any of its declarations are read. 

27 """ 

28 

29 def __init__(self) -> None: 

30 super().__init__( 

31 target=TreeBuilder(), 

32 forbid_dtd=True, 

33 forbid_entities=True, 

34 forbid_external=True, 

35 ) 

36 

37 def defused_start_doctype_decl( 

38 self, name, sysid, pubid, has_internal_subset 

39 ): 

40 if has_internal_subset: 

41 raise DTDForbidden(name, sysid, pubid) 

42 

43 

44class EverNoteLoader(BaseLoader): 

45 """Extract notes from an ENEX file without an undeclared parser dependency.""" 

46 

47 def __init__( 

48 self, file_path: str | Path, load_single_document: bool = True 

49 ) -> None: 

50 self.file_path = str(file_path) 

51 self.load_single_document = load_single_document 

52 

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

54 notes = self._iter_notes() 

55 if self.load_single_document: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true

56 yield Document( 

57 page_content="".join(note.page_content for note in notes), 

58 metadata={"source": self.file_path}, 

59 ) 

60 else: 

61 yield from notes 

62 

63 def _iter_notes(self) -> Iterator[Document]: 

64 # Standard exports declare Evernote's doctype by its external 

65 # identifier only, and ElementTree does not fetch that DTD. Any 

66 # internal subset is refused (see _EvernoteXMLParser), and entity 

67 # declarations and external references stay forbidden. Clearing the 

68 # root after each note avoids retaining earlier notes and attachments. 

69 with open(self.file_path, "rb") as stream: 

70 context = iterparse( 

71 stream, 

72 events=("start", "end"), 

73 parser=_EvernoteXMLParser(), 

74 ) 

75 root = None 

76 for event, element in context: 

77 if root is None: 

78 root = element 

79 if event != "end" or element.tag != "note": 

80 continue 

81 

82 content = element.findtext("content") 

83 if content is not None: 83 ↛ 109line 83 didn't jump to line 109 because the condition on line 83 was always true

84 metadata: dict[str, Any] = {"source": self.file_path} 

85 for child in element: 

86 if child.tag in {"content", "resource"}: 

87 continue 

88 if child.tag == "note-attributes": 

89 for attribute in child: 

90 metadata[f"note-attributes.{attribute.tag}"] = ( 

91 attribute.text 

92 ) 

93 elif child.tag in {"created", "updated"}: 

94 metadata[child.tag] = ( 

95 strptime(child.text, "%Y%m%dT%H%M%SZ") 

96 if child.text 

97 else None 

98 ) 

99 else: 

100 metadata[child.tag] = child.text 

101 

102 html = BeautifulSoup(content, "html.parser") 

103 for ignored in html(["script", "style"]): 

104 ignored.decompose() 

105 yield Document( 

106 page_content=html.get_text(" ", strip=True), 

107 metadata=metadata, 

108 ) 

109 element.clear() 

110 root.clear()