From a8a7b9e7b88285400651abdc1b5db3d10fb0d2a7 Mon Sep 17 00:00:00 2001 From: Connor Olding Date: Fri, 4 Sep 2020 15:07:48 +0200 Subject: [PATCH] deduplicate main loop --- respodns/dns.py | 38 ++++++++++++++++++-------------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/respodns/dns.py b/respodns/dns.py index cf78378..bfbfc94 100644 --- a/respodns/dns.py +++ b/respodns/dns.py @@ -215,38 +215,36 @@ async def main(db, filepaths, checks, opts: Options): pooler = make_pooler(opts.ip_simul, finisher) + seen = 0 + async def process(ip, total=None): + nonlocal seen + first = seen == 0 + seen += 1 + if opts.progress: + if total is None: + print(f"#{seen}: {ip}", file=stderr) + else: + print(f"#{seen}/{total}: {ip}", file=stderr) + stderr.flush() + if not first: + await sleep(opts.ip_wait) + await pooler(try_ip(db, ip, checks, opts)) + if opts.blocking_file_io: from .ip_util import read_ips - - seen = 0 for filepath in filepaths: f = stdin if filepath == "" else open(filepath, "r") for ip in read_ips(f): - first = seen == 0 - seen += 1 - if opts.progress: - print(f"#{seen}: {ip}", file=stderr) - stderr.flush() - if not first: - await sleep(opts.ip_wait) - await pooler(try_ip(db, ip, checks, opts)) + await process(ip) if f != stdin: f.close() - else: # blocking_file_io + else: from .ip_util import IpReader fps = [stdin if fp == "" else fp for fp in filepaths] with IpReader(*fps) as reader: - seen = 0 async for ip in reader: - first = seen == 0 - seen += 1 - if opts.progress: - print(f"#{seen}/{reader.total}: {ip}", file=stderr) - stderr.flush() - if not first: - await sleep(opts.ip_wait) - await pooler(try_ip(db, ip, checks, opts)) + await process(ip, reader.total) await pooler() await syncing