attachments.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. # SPDX-FileCopyrightText: 2026 geisserml <geisserml@gmail.com>
  2. # SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
  3. from pathlib import Path
  4. from pypdfium2._cli._parsers import (
  5. add_input, get_input,
  6. parse_numtext,
  7. )
  8. ACTION_LIST = "list"
  9. ACTION_EXTRACT = "extract"
  10. ACTION_EDIT = "edit"
  11. def attach(parser): # hook
  12. add_input(parser, pages=False)
  13. subparsers = parser.add_subparsers(dest="action")
  14. subparsers.add_parser(ACTION_LIST)
  15. parser_extract = subparsers.add_parser(ACTION_EXTRACT)
  16. parser_extract.add_argument(
  17. "--numbers",
  18. type = parse_numtext,
  19. )
  20. parser_extract.add_argument(
  21. "--output-dir", "-o",
  22. type = Path,
  23. required = True,
  24. )
  25. parser_edit = subparsers.add_parser(ACTION_EDIT)
  26. parser_edit.add_argument(
  27. "--del-numbers", "-d",
  28. type = parse_numtext,
  29. )
  30. parser_edit.add_argument(
  31. "--add-files", "-a",
  32. nargs = "+",
  33. metavar = "F",
  34. type = Path,
  35. )
  36. parser_edit.add_argument(
  37. "--output", "-o",
  38. type = Path,
  39. required = True,
  40. )
  41. def main(args):
  42. pdf = get_input(args)
  43. n_attachments = pdf.count_attachments()
  44. if args.action == ACTION_LIST:
  45. for i in range(n_attachments):
  46. attachment = pdf.get_attachment(i)
  47. print(f"[{i+1}]", attachment.get_name())
  48. elif args.action == ACTION_EXTRACT:
  49. if not args.numbers:
  50. args.numbers = range(n_attachments)
  51. n_digits = len(str( max(args.numbers) + 1 ))
  52. for i in args.numbers:
  53. attachment = pdf.get_attachment(i)
  54. name = attachment.get_name()
  55. out_path = args.output_dir / ("%0*d_%s" % (n_digits, i+1, name))
  56. out_path.write_bytes( attachment.get_data() )
  57. elif args.action == ACTION_EDIT:
  58. if args.del_numbers:
  59. for i in sorted(args.del_numbers, reverse=True):
  60. pdf.del_attachment(i)
  61. if args.add_files:
  62. for fp in args.add_files:
  63. attachment = pdf.new_attachment(fp.name)
  64. attachment.set_data( fp.read_bytes() )
  65. pdf.save(args.output)
  66. else:
  67. assert False