refactor bot into boot + book client

bookclient.py

@@ -0,0 +1,106 @@
+from threading import Thread
+from ircclient import IRCClient, MODE_SEARCH, MODE_BOOK
+import logging
+import queue
+import subprocess
+
+logging.basicConfig(level=logging.DEBUG)
+log = logging.getLogger(__name__)
+log.setLevel(logging.DEBUG)
+
+def get_books_from_list(filename):
+    f = open(filename, "r")
+    ret = []
+    for line in f.readlines():
+        line = line.lower().strip()
+        if not line.startswith('!') or not any(_type in line for _type in ['epub', 'mobi']):
+            continue
+        ret.append(line)
+        log.info("Book matches: %s", line)
+    ret = list(set(ret)) # dedup
+    return ret
+
+def mode_from_files(files):
+    for f in files:
+        f = f.lower()
+        if 'searchbot' in f or 'searchook' in f.lower():
+            return MODE_SEARCH
+    return MODE_BOOK
+
+def handle_files(files):
+    log.info("Unarchived files %s", files)
+    out = []
+    mode = mode_from_files(files)
+    if mode == MODE_SEARCH:
+        for f in files:
+            if "searchbot" not in f.lower() and "searchook" not in f.lower():
+                continue
+            out.extend(get_books_from_list(f))
+        return out
+
+    for f in files:
+        if f.lower().endswith(".epub"):
+            log.info("EPUB %s", f)
+            new_fname = f.replace("epub", "mobi")
+            p = subprocess.Popen(["ebook-convert", f, new_fname], stdout=subprocess.DEVNULL)
+            # TODO log to file?
+            p.wait()
+            out.append(new_fname)
+        out.append(f)
+    # TODO make paths absolute?
+    return out
+
+def handle_results(q, cb):
+    while True:
+        try:
+            item = q.get(timeout=1)
+            if item is None:
+                break
+        except queue.Empty:
+            continue
+        log.info("Got a result! %s", item)
+        if item['type'] == 'status':
+            cb((item['key'], item['status']))
+        elif item['type'] == 'files':
+            cb(handle_files(item['files']))
+
+class BookClient:
+    workers = []
+    def __init__(self, cb):
+        self.q = queue.Queue()
+        self.rq = queue.Queue()
+
+        self.results_t = Thread(target=handle_results, args=(self.rq, cb))
+        self.results_t.daemon = True
+        self.results_t.start()
+
+        self.create_worker()
+
+    def create_worker(self):
+        print("Spawning a new worker")
+        worker = IRCClient(command_queue=self.q, results_queue=self.rq)
+        worker.start()
+        self.workers.append(worker)
+
+    def free_workers(self):
+        return any([not worker.busy for worker in self.workers])
+
+    def request(self, mode, query):
+        if not self.free_workers():
+            self.create_worker()
+
+        if mode not in [MODE_SEARCH, MODE_BOOK]:
+            print('Invalid mode')
+            return
+
+        self.q.put({'query': query, 'mode': mode})
+        print("Command acknowledged")
+
+    def __del__(self):
+        print("Cleaning up workers")
+        for worker in self.workers:
+            worker.stop()
+            worker.join(timeout=5)
+        print("Cleaning up results")
+        self.rq.put(None)
+        self.results_t.join(timeout=5)

bot.py

@@ -1,18 +1,12 @@
 #!/usr/bin/env python3.7
 import logging
-import queue
 import shlex
-import subprocess
-import sys
 import readline
 import os
+from collections.abc import Iterable
 
-from threading import Thread
-from ircclient import IRCClient, MODE_SEARCH, MODE_BOOK
+from bookclient import BookClient
 
-logging.basicConfig(level=logging.DEBUG)
-log = logging.getLogger(__name__)
-log.setLevel(logging.DEBUG)
 histfile = os.path.join(os.path.expanduser("~"), ".book_history")
 
 def usage():
@@ -20,70 +14,16 @@ def usage():
     print("\t SEARCH <BOOK>")
     print("\t BOOK <BOT COMMAND>")
 
-def get_books_from_list(filename):
-    f = open(filename, "r")
-    ret = []
-    for line in f.readlines():
-        line = line.lower().strip()
-        if not line.startswith('!') or not any(_type in line for _type in ['epub', 'mobi']):
-            continue
-        ret.append(line)
-        log.info("Book matches: %s", line)
-    ret = list(set(ret)) # dedup
-    return ret
-
-def mode_from_files(files):
-    for f in files:
-        f = f.lower()
-        if 'searchbot' in f or 'searchook' in f.lower():
-            return MODE_SEARCH
-    return MODE_BOOK
-
-def handle_files(files):
-    log.info("Unarchived files %s", files)
-    out = []
-    mode = mode_from_files(files)
-    if mode == MODE_SEARCH:
-        for f in files:
-            if "searchbot" not in f.lower() and "searchook" not in f.lower():
-                continue
-            out.extend(get_books_from_list(f))
-        return
-
-    for f in files:
-        if f.lower().endswith(".epub"):
-            log.info("EPUB %s", f)
-            new_fname = f.replace("epub", "mobi")
-            p = subprocess.Popen(["ebook-convert", f, new_fname], stdout=subprocess.DEVNULL)
-            # TODO log to file?
-            p.wait()
-            out.append(new_fname)
-        out.append(f)
-    # TODO make paths absolute?
-    for filename in out:
-        log.info(filename)
-
-def handle_results(q):
-    while True:
-        item = q.get()
-        log.info("Got a result! %s", item)
-        if item['type'] == 'status':
-            print(item['key'], item['status'])
-        elif item['type'] == 'files':
-            handle_files(item['files'])
+def cb(arg):
+    print("Called callback!")
+    if isinstance(arg, Iterable):
+        for item in arg:
+            print(item)
+    else:
+        print(arg)
 
 def main():
-    q = queue.Queue()
-    rq = queue.Queue()
-
-    workers = []
-    worker = IRCClient(command_queue=q, results_queue=rq)
-    worker.start()
-    workers.append(worker)
-
-    results_t = Thread(target=handle_results, args=(rq,))
-    results_t.daemon = True
-    results_t.start()
+    b = BookClient(cb)
     while True:
         line = input('> ').strip().lower()
         try:
@@ -91,28 +31,14 @@ def main():
         except Exception as e:
             print(e)
             continue
-        if len(split) == 0:
-            continue
 
-        if split[0].lower() not in [MODE_SEARCH, MODE_BOOK]:
-            usage()
+        if len(split) == 0:
             continue
 
-        if all([ worker.busy for worker in workers]):
-            print("Spawning a new worker as all are busy...")
-            worker = IRCClient(command_queue=q, results_queue=rq)
-            worker.start()
-            workers.append(worker)
-
         mode = split[0]
         query = " ".join(split[1:])
-        data = {'query': query, 'mode': mode}
-        q.put(data)
-        print("Command acknowledged")
+        b.request(mode, query)
 
-    results_t.join()
-    for worker in workers:
-        worker.join()
 
 if __name__ == "__main__":
     try:

ircclient.py

@@ -12,8 +12,6 @@ from threading import Thread
 from collections import defaultdict
 
 logging.basicConfig(level=logging.DEBUG)
-log = logging.getLogger(__name__)
-log.setLevel(logging.DEBUG)
 
 MODE_SEARCH = 'search'
 MODE_BOOK = 'book'
@@ -81,6 +79,7 @@ class IRCClient(Thread):
     time_joined = None
     jobs = defaultdict(dict)
     busy = False
+    running = True
 
     def __init__(self, command_queue, results_queue):
         super(IRCClient, self).__init__(daemon=True)
@@ -89,15 +88,20 @@ class IRCClient(Thread):
         self.results_queue = results_queue
         self.send_queue = queue.Queue()
 
+        self.log = logging.getLogger(self.getName())
+        self.log.setLevel(logging.DEBUG)
+
         nickstr = "NICK %s" % self.name
         userstr = "USER %s %s bla :%s" % (self.name, self.HOST, self.name)
 
         self.send_queue.put(nickstr)
         self.send_queue.put(userstr)
 
+    def stop(self):
+        self.running = False
     def run(self):
         self.handle_connect()
-        while True:
+        while self.running:
             if self.connected and self.joined_channel:
                 self.handle_commands()
             self.process_send_queue()
@@ -106,14 +110,14 @@ class IRCClient(Thread):
             except socket.timeout:
                 continue
             except socket.error as e:
-                log.error('socket error')
-                log.exception(e)
+                self.log.error('socket error')
+                self.log.exception(e)
                 self.handle_connect()
             except Exception as e:
-                log.exception(e)
+                self.log.exception(e)
                 break
         self.handle_close()
-        log.info('Exiting RUN')
+        self.log.info('Exiting RUN')
 
     def handle_commands(self):
         if self.command_queue.empty():
@@ -122,21 +126,24 @@ class IRCClient(Thread):
 
         elapsed = time.time() - self.time_joined
         if elapsed < self.TIME_TO_FIRST_COMMAND:
-            log.info("commands to process, but we have to wait %d", self.TIME_TO_FIRST_COMMAND - elapsed)
+            self.log.info("commands to process, but we have to wait %d seconds", self.TIME_TO_FIRST_COMMAND - elapsed)
             time.sleep(1)
             return
         command = self.command_queue.get()
-        log.info("command %s", command)
+        self.log.info("command %s", command)
         job = query_to_job_key(command['query'])
         self.set_job_state(job, 'pending')
 
         if command['mode'] == MODE_SEARCH:
             self.send_queue.put("PRIVMSG %s :@%s %s " % (self.CHANNEL, self.SEARCH_BOT, command['query']))
             return
-        self.send_queue.put("PRIVMSG %s :%s " % (self.CHANNEL, command['query']))
+        elif command['mode'] == MODE_BOOK:
+            self.send_queue.put("PRIVMSG %s :%s " % (self.CHANNEL, command['query']))
+        else:
+            self.log.error('Invalid command')
 
     def set_job_state(self, job, state):
-        log.info("Setting job [%s] to %s", job, state)
+        self.log.info("Setting job [%s] to %s", job, state)
         self.jobs[job].update({'state': state})
         self.results_queue.put({'type': 'status', 'status': state, 'key': job})
         if state == 'done':
@@ -146,11 +153,11 @@ class IRCClient(Thread):
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.socket.connect((self.HOST, self.PORT))
         self.socket.settimeout(2)
-        log.info("connected")
+        self.log.info("connected")
         self.connected = True
 
     def handle_close(self):
-        log.info("closed")
+        self.log.info("closed")
         self.connected = False
         self.socket.close()
 
@@ -181,7 +188,7 @@ class IRCClient(Thread):
             comm = words[1]
 
             if self.joined_channel:
-                # log.debug(line)
+                # self.log.debug(line)
                 pass
 
             if comm in self.IGNORE:
@@ -190,7 +197,7 @@ class IRCClient(Thread):
             if comm == "JOIN":
                 if self.name not in msg_from:  # msg "NICK joined the channel" not about me
                     continue
-                log.info("Joined channel %s", self.CHANNEL)
+                self.log.info("Joined channel %s", self.CHANNEL)
                 self.time_joined = time.time()
                 self.joined_channel = True
                 continue
@@ -199,7 +206,7 @@ class IRCClient(Thread):
                 # private message not addressed to me
                 if words[2] != self.name:
                     continue
-                log.info("privmsg: %s", line)
+                self.log.info("privmsg: %s", line)
                 dcc_args = get_dcc_args(line)
                 if dcc_args is None:
                     continue
@@ -226,9 +233,9 @@ class IRCClient(Thread):
 
     def netcat(self, ip, port, size, filename, job_key):
         filename = os.path.basename(filename).replace(" ", "_")
-        log.info('netcat: %s %d %d %s', ip, port, size, filename)
+        self.log.info('netcat: %s %d %d %s', ip, port, size, filename)
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        log.info("Receiving file")
+        self.log.info("Receiving file")
         s.connect((ip, port))
 
         fname = os.path.join(self.PATH, filename)
@@ -238,13 +245,13 @@ class IRCClient(Thread):
         while True:
             data = s.recv(16384)
             if len(data) == 0:
-                log.info("No data received - finished")
+                self.log.info("No data received - finished")
                 break
             count += len(data)
             f.write(data)
             perc = int(100 * count / size)
             if perc % 10 == 0 and perc != last_perc:
-                log.info("Download percentage: %d", perc)
+                self.log.info("Download percentage: %d", perc)
                 self.set_job_state(job_key, "%d%%" % perc)
                 last_perc = perc
             if count >= size:
@@ -265,6 +272,6 @@ class IRCClient(Thread):
             return
         data = self.send_queue.get()
         if not data.startswith("PONG"):
-            log.info("Sending %s", data)
+            self.log.info("Sending %s", data)
         add = bytes(str(data), "utf-8") + bytes([13, 10])
         self.socket.send(add)