From a8a7b9e7b88285400651abdc1b5db3d10fb0d2a7 Mon Sep 17 00:00:00 2001
From: Connor Olding <cloningdonor@gmail.com>
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