helpers.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. """Document loader helpers."""
  2. import concurrent.futures
  3. from typing import NamedTuple, Optional, cast
  4. class FileEncoding(NamedTuple):
  5. """A file encoding as the NamedTuple."""
  6. encoding: Optional[str]
  7. """The encoding of the file."""
  8. confidence: float
  9. """The confidence of the encoding."""
  10. language: Optional[str]
  11. """The language of the file."""
  12. def detect_file_encodings(file_path: str, timeout: int = 5) -> list[FileEncoding]:
  13. """Try to detect the file encoding.
  14. Returns a list of `FileEncoding` tuples with the detected encodings ordered
  15. by confidence.
  16. Args:
  17. file_path: The path to the file to detect the encoding for.
  18. timeout: The timeout in seconds for the encoding detection.
  19. """
  20. import chardet
  21. def read_and_detect(file_path: str) -> list[dict]:
  22. with open(file_path, "rb") as f:
  23. rawdata = f.read()
  24. return cast(list[dict], chardet.detect_all(rawdata))
  25. with concurrent.futures.ThreadPoolExecutor() as executor:
  26. future = executor.submit(read_and_detect, file_path)
  27. try:
  28. encodings = future.result(timeout=timeout)
  29. except concurrent.futures.TimeoutError:
  30. raise TimeoutError(
  31. f"Timeout reached while detecting encoding for {file_path}"
  32. )
  33. if all(encoding["encoding"] is None for encoding in encodings):
  34. raise RuntimeError(f"Could not detect encoding for {file_path}")
  35. return [FileEncoding(**enc) for enc in encodings if enc["encoding"] is not None]