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

1""" 

2Read text uploads in the encoding they were written in. 

3 

4LangChain's ``autodetect_encoding`` passes every candidate that 

5``chardet.detect_all`` returns to ``FileEncoding(**candidate)``. chardet 7 

6adds a ``mime_type`` key to each candidate, so detection raises ``TypeError`` 

7and every ``.txt`` file that is not UTF-8 is refused, readable or not. With an 

8older chardet, the same detection decodes binary bytes into text. Which of the 

9two happened depended on the installed chardet (#6532). 

10 

11These loaders decide the encoding themselves, the same way everywhere. The 

12scope is deliberately narrow: UTF-8, UTF-16 and UTF-32 with a byte-order mark, 

13and Windows-1252, the Western-European legacy encoding Notepad and Excel write. 

14A file in any other legacy encoding has one of two outcomes: 

15 

16- it is refused, when its bytes include one Windows-1252 leaves undefined or 

17 decode into control characters (cp866, for one); 

18- otherwise it is read as Windows-1252 and stored as mojibake (KOI8-R, 

19 Shift_JIS, GB18030, EUC-KR and EUC-JP text all decode this way). 

20 

21The encoding each file was read in is recorded on its documents as 

22``detected_encoding``, for callers that keep the documents. The upload path 

23joins their text in ``extract_text_from_bytes`` and does not keep it yet 

24(#6836). 

25""" 

26 

27import re 

28from collections.abc import Iterator 

29from pathlib import Path 

30 

31from langchain_community.document_loaders import CSVLoader, TextLoader 

32from langchain_core.documents import Document 

33 

34# Characters that legacy text does not contain: NUL and the C0 controls other 

35# than tab, line feed, vertical tab, form feed, carriage return, the DOS 

36# end-of-file mark (SUB) and escape. 

37_NOT_TEXT = re.compile(r"[\x00-\x08\x0e-\x19\x1c-\x1f]") 

38 

39 

40def _fallback_encodings(raw: bytes) -> list[str]: 

41 # UTF-32's little-endian mark begins with UTF-16's. 

42 if raw.startswith((b"\xff\xfe\x00\x00", b"\x00\x00\xfe\xff")): 

43 return ["utf-32", "cp1252"] 

44 if raw.startswith((b"\xff\xfe", b"\xfe\xff")): 

45 return ["utf-16", "cp1252"] 

46 # What Notepad and Excel write on Western-European Windows ("ANSI"). 

47 return ["cp1252"] 

48 

49 

50def detect_text_encoding(raw: bytes) -> str | None: 

51 """ 

52 Return the encoding that decodes ``raw`` into text, or None. 

53 

54 Bytes that are UTF-8, with or without a byte-order mark, are read as 

55 UTF-8, as before. Other bytes are read as UTF-16 or UTF-32 when they start 

56 with its byte-order mark, and otherwise as Windows-1252, but only when the 

57 result is text: bytes that decode into control characters are refused. 

58 """ 

59 try: 

60 raw.decode("utf-8-sig") 

61 return "utf-8-sig" 

62 except UnicodeDecodeError: 

63 pass 

64 for encoding in _fallback_encodings(raw): 

65 try: 

66 text = raw.decode(encoding) 

67 except UnicodeDecodeError: 

68 continue 

69 if not _NOT_TEXT.search(text): 

70 return encoding 

71 return None 

72 

73 

74def _file_encoding(file_path: str | Path) -> str: 

75 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Encoding loader reads its explicit local file argument; byte uploads pass a server-created temporary file 

76 encoding = detect_text_encoding(Path(file_path).read_bytes()) 

77 if encoding is None: 

78 raise ValueError( 

79 f"{Path(file_path).name}: no recognized text encoding (tried " 

80 "UTF-8, UTF-16 or UTF-32 with a byte-order mark, and " 

81 "Windows-1252)" 

82 ) 

83 return encoding 

84 

85 

86def _with_encoding( 

87 documents: Iterator[Document], encoding: str 

88) -> Iterator[Document]: 

89 for document in documents: 

90 document.metadata["detected_encoding"] = encoding 

91 yield document 

92 

93 

94class EncodingAwareTextLoader(TextLoader): 

95 """A ``TextLoader`` that reads a file in the encoding it was written in.""" 

96 

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

98 self.encoding = _file_encoding(self.file_path) 

99 yield from _with_encoding(super().lazy_load(), self.encoding) 

100 

101 

102class EncodingAwareCSVLoader(CSVLoader): 

103 """A ``CSVLoader`` that reads a file in the encoding it was written in. 

104 

105 A byte-order mark is dropped rather than read into the first column name. 

106 """ 

107 

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

109 self.encoding = _file_encoding(self.file_path) 

110 yield from _with_encoding(super().lazy_load(), self.encoding)