base.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. from loguru import logger
  2. from magic_pdf.libs.commons import read_file
  3. from magic_pdf.libs.drop_reason import DropReason
  4. from magic_pdf.spark.s3 import get_s3_config
  5. def get_data_source(jso: dict):
  6. data_source = jso.get("data_source")
  7. if data_source is None:
  8. data_source = jso.get("file_source")
  9. return data_source
  10. def get_data_type(jso: dict):
  11. data_type = jso.get("data_type")
  12. if data_type is None:
  13. data_type = jso.get("file_type")
  14. return data_type
  15. def get_bookid(jso: dict):
  16. book_id = jso.get("bookid")
  17. if book_id is None:
  18. book_id = jso.get("original_file_id")
  19. return book_id
  20. def exception_handler(jso: dict, e):
  21. logger.exception(e)
  22. jso["need_drop"] = True
  23. jso["drop_reason"] = DropReason.Exception
  24. jso["exception"] = f"ERROR: {e}"
  25. return jso
  26. def get_bookname(jso: dict):
  27. data_source = get_data_source(jso)
  28. file_id = jso.get("file_id")
  29. book_name = f"{data_source}/{file_id}"
  30. return book_name
  31. def get_pdf_bytes(jso: dict):
  32. pdf_s3_path = jso.get("file_location")
  33. s3_config = get_s3_config(pdf_s3_path)
  34. pdf_bytes = read_file(pdf_s3_path, s3_config)
  35. return pdf_bytes