From 05bd7cbcc62588e431c5efd533189e246b2a997e Mon Sep 17 00:00:00 2001 From: Bryan Newbold Date: Tue, 26 Oct 2021 12:54:37 -0700 Subject: make fmt --- python/pdftrio_tool.py | 82 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 50 insertions(+), 32 deletions(-) (limited to 'python/pdftrio_tool.py') diff --git a/python/pdftrio_tool.py b/python/pdftrio_tool.py index e195bc7..9316313 100755 --- a/python/pdftrio_tool.py +++ b/python/pdftrio_tool.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 - """ Basically just a copy of grobid_tool.py, but for PDF classification instead of text extraction. @@ -21,19 +20,29 @@ def run_classify_pdf_json(args): pdftrio_client = PdfTrioClient(host_url=args.pdftrio_host) wayback_client = WaybackClient() if args.jobs > 1: - worker = PdfTrioWorker(pdftrio_client, wayback_client, sink=None, mode=args.pdftrio_mode) + worker = PdfTrioWorker(pdftrio_client, + wayback_client, + sink=None, + mode=args.pdftrio_mode) multi_worker = MultiprocessWrapper(worker, args.sink) pusher = JsonLinePusher(multi_worker, args.json_file, batch_size=args.jobs) else: - worker = PdfTrioWorker(pdftrio_client, wayback_client, sink=args.sink, mode=args.pdftrio_mode) + worker = PdfTrioWorker(pdftrio_client, + wayback_client, + sink=args.sink, + mode=args.pdftrio_mode) pusher = JsonLinePusher(worker, args.json_file) pusher.run() + def run_classify_pdf_cdx(args): pdftrio_client = PdfTrioClient(host_url=args.pdftrio_host) wayback_client = WaybackClient() if args.jobs > 1: - worker = PdfTrioWorker(pdftrio_client, wayback_client, sink=None, mode=args.pdftrio_mode) + worker = PdfTrioWorker(pdftrio_client, + wayback_client, + sink=None, + mode=args.pdftrio_mode) multi_worker = MultiprocessWrapper(worker, args.sink) pusher = CdxLinePusher( multi_worker, @@ -43,7 +52,10 @@ def run_classify_pdf_cdx(args): batch_size=args.jobs, ) else: - worker = PdfTrioWorker(pdftrio_client, wayback_client, sink=args.sink, mode=args.pdftrio_mode) + worker = PdfTrioWorker(pdftrio_client, + wayback_client, + sink=args.sink, + mode=args.pdftrio_mode) pusher = CdxLinePusher( worker, args.cdx_file, @@ -52,6 +64,7 @@ def run_classify_pdf_cdx(args): ) pusher.run() + def run_classify_pdf_zipfile(args): pdftrio_client = PdfTrioClient(host_url=args.pdftrio_host) worker = PdfTrioBlobWorker(pdftrio_client, sink=args.sink, mode=args.pdftrio_mode) @@ -60,48 +73,53 @@ def run_classify_pdf_zipfile(args): 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") parser.add_argument('--pdftrio-host', - default="http://pdftrio.qa.fatcat.wiki", - help="pdftrio API host/port") + default="http://pdftrio.qa.fatcat.wiki", + help="pdftrio API host/port") parser.add_argument('--pdftrio-mode', - default="auto", - help="which classification mode to use") + default="auto", + help="which classification mode to use") subparsers = parser.add_subparsers() - sub_classify_pdf_json = subparsers.add_parser('classify-pdf-json', + sub_classify_pdf_json = subparsers.add_parser( + 'classify-pdf-json', help="for each JSON line with CDX info, fetches PDF and does pdftrio classify_pdfion") sub_classify_pdf_json.set_defaults(func=run_classify_pdf_json) sub_classify_pdf_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_classify_pdf_cdx = subparsers.add_parser('classify-pdf-cdx', + sub_classify_pdf_cdx = subparsers.add_parser( + 'classify-pdf-cdx', help="for each CDX line, fetches PDF and does pdftrio classify_pdfion") sub_classify_pdf_cdx.set_defaults(func=run_classify_pdf_cdx) sub_classify_pdf_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_classify_pdf_zipfile = subparsers.add_parser('classify-pdf-zipfile', - help="opens zipfile, iterates over PDF files inside and does pdftrio classify_pdf for each") + sub_classify_pdf_zipfile = subparsers.add_parser( + 'classify-pdf-zipfile', + help= + "opens zipfile, iterates over PDF files inside and does pdftrio classify_pdf for each") sub_classify_pdf_zipfile.set_defaults(func=run_classify_pdf_zipfile) sub_classify_pdf_zipfile.add_argument('zip_file', - help="zipfile with PDFs to classify", - type=str) + help="zipfile with PDFs to classify", + type=str) args = parser.parse_args() if not args.__dict__.get("func"): @@ -112,10 +130,10 @@ def main(): if args.kafka_mode: produce_topic = "sandcrawler-{}.pdftrio-output".format(args.kafka_env) print("Running in kafka output mode, publishing to {}\n".format(produce_topic)) - args.sink = KafkaSink(kafka_hosts=args.kafka_hosts, - produce_topic=produce_topic) + args.sink = KafkaSink(kafka_hosts=args.kafka_hosts, produce_topic=produce_topic) args.func(args) + if __name__ == '__main__': main() -- cgit v1.2.3