diff options
Diffstat (limited to 'python/pdfextract_tool.py')
-rwxr-xr-x | python/pdfextract_tool.py | 74 |
1 files changed, 39 insertions, 35 deletions
diff --git a/python/pdfextract_tool.py b/python/pdfextract_tool.py index 89ecf1c..717b743 100755 --- a/python/pdfextract_tool.py +++ b/python/pdfextract_tool.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 - """ KNOWN ISSUE: thumbnails are not published to kafka in multi-processing mode """ @@ -20,10 +19,13 @@ def run_extract_json(args): multi_worker = MultiprocessWrapper(worker, args.sink) pusher = JsonLinePusher(multi_worker, args.json_file, batch_size=args.jobs) else: - worker = PdfExtractWorker(wayback_client, sink=args.sink, thumbnail_sink=args.thumbnail_sink) + worker = PdfExtractWorker(wayback_client, + sink=args.sink, + thumbnail_sink=args.thumbnail_sink) pusher = JsonLinePusher(worker, args.json_file) pusher.run() + def run_extract_cdx(args): wayback_client = WaybackClient() if args.jobs > 1: @@ -37,7 +39,9 @@ def run_extract_cdx(args): batch_size=args.jobs, ) else: - worker = PdfExtractWorker(wayback_client, sink=args.sink, thumbnail_sink=args.thumbnail_sink) + worker = PdfExtractWorker(wayback_client, + sink=args.sink, + thumbnail_sink=args.thumbnail_sink) pusher = CdxLinePusher( worker, args.cdx_file, @@ -46,6 +50,7 @@ def run_extract_cdx(args): ) pusher.run() + def run_extract_zipfile(args): if args.jobs > 1: print("multi-processing: {}".format(args.jobs), file=sys.stderr) @@ -57,6 +62,7 @@ def run_extract_zipfile(args): pusher = ZipfilePusher(worker, args.zip_file) pusher.run() + def run_single(args): worker = PdfExtractBlobWorker(sink=args.sink, thumbnail_sink=args.thumbnail_sink) with open(args.pdf_file, 'rb') as pdf_file: @@ -67,51 +73,48 @@ def run_single(args): args.thumbnail_sink.finish() - def main(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--kafka-mode', - action='store_true', - help="send output to Kafka (not stdout)") + action='store_true', + help="send output to Kafka (not stdout)") parser.add_argument('--kafka-hosts', - default="localhost:9092", - help="list of Kafka brokers (host/port) to use") + default="localhost:9092", + help="list of Kafka brokers (host/port) to use") parser.add_argument('--kafka-env', - default="dev", - help="Kafka topic namespace to use (eg, prod, qa, dev)") - parser.add_argument('-j', '--jobs', - default=8, type=int, - help="parallelism for batch CPU jobs") + default="dev", + help="Kafka topic namespace to use (eg, prod, qa, dev)") + parser.add_argument('-j', + '--jobs', + default=8, + type=int, + help="parallelism for batch CPU jobs") subparsers = parser.add_subparsers() - sub_extract_json = subparsers.add_parser('extract-json', + sub_extract_json = subparsers.add_parser( + 'extract-json', help="for each JSON line with CDX info, fetches PDF and does PDF extraction") sub_extract_json.set_defaults(func=run_extract_json) sub_extract_json.add_argument('json_file', - help="JSON file to import from (or '-' for stdin)", - type=argparse.FileType('r')) + help="JSON file to import from (or '-' for stdin)", + type=argparse.FileType('r')) - sub_extract_cdx = subparsers.add_parser('extract-cdx', - help="for each CDX line, fetches PDF and does PDF extraction") + sub_extract_cdx = subparsers.add_parser( + 'extract-cdx', help="for each CDX line, fetches PDF and does PDF extraction") sub_extract_cdx.set_defaults(func=run_extract_cdx) sub_extract_cdx.add_argument('cdx_file', - help="CDX file to import from (or '-' for stdin)", - type=argparse.FileType('r')) + help="CDX file to import from (or '-' for stdin)", + type=argparse.FileType('r')) - sub_extract_zipfile = subparsers.add_parser('extract-zipfile', + sub_extract_zipfile = subparsers.add_parser( + 'extract-zipfile', help="opens zipfile, iterates over PDF files inside and does PDF extract for each") sub_extract_zipfile.set_defaults(func=run_extract_zipfile) - sub_extract_zipfile.add_argument('zip_file', - help="zipfile with PDFs to extract", - type=str) + sub_extract_zipfile.add_argument('zip_file', help="zipfile with PDFs to extract", type=str) - sub_single = subparsers.add_parser('single', - help="opens single PDF and extracts it") + sub_single = subparsers.add_parser('single', help="opens single PDF and extracts it") sub_single.set_defaults(func=run_single) - sub_single.add_argument('pdf_file', - help="single PDF to extract", - type=str) + sub_single.add_argument('pdf_file', help="single PDF to extract", type=str) args = parser.parse_args() if not args.__dict__.get("func"): @@ -123,17 +126,18 @@ def main(): if args.kafka_mode: text_topic = "sandcrawler-{}.pdf-text".format(args.kafka_env) thumbnail_topic = "sandcrawler-{}.pdf-thumbnail-180px-jpg".format(args.kafka_env) - args.sink = KafkaCompressSink(kafka_hosts=args.kafka_hosts, - produce_topic=text_topic) + args.sink = KafkaCompressSink(kafka_hosts=args.kafka_hosts, produce_topic=text_topic) args.thumbnail_sink = KafkaSink(kafka_hosts=args.kafka_hosts, - produce_topic=thumbnail_topic) + produce_topic=thumbnail_topic) print("Running in kafka output mode, publishing to {} and {}\n".format( - text_topic, thumbnail_topic), file=sys.stderr) + text_topic, thumbnail_topic), + file=sys.stderr) else: args.sink = None args.thumbnail_sink = None args.func(args) + if __name__ == '__main__': main() |