bookclient.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. from threading import Thread
  2. from ircclient import IRCClient, MODE_SEARCH, MODE_BOOK
  3. import logging
  4. import queue
  5. import subprocess
  6. logging.basicConfig(level=logging.DEBUG)
  7. log = logging.getLogger(__name__)
  8. log.setLevel(logging.DEBUG)
  9. def get_books_from_list(filename):
  10. f = open(filename, "r")
  11. ret = []
  12. for line in f.readlines():
  13. line = line.lower().strip()
  14. if not line.startswith('!') or not any(_type in line for _type in ['epub', 'mobi']):
  15. continue
  16. ret.append(line)
  17. log.info("Book matches: %s", line)
  18. ret = list(set(ret)) # dedup
  19. return ret
  20. def mode_from_files(files):
  21. for f in files:
  22. f = f.lower()
  23. if 'searchbot' in f or 'searchook' in f.lower():
  24. return MODE_SEARCH
  25. return MODE_BOOK
  26. def handle_files(files):
  27. log.info("Unarchived files %s", files)
  28. out = []
  29. mode = mode_from_files(files)
  30. if mode == MODE_SEARCH:
  31. for f in files:
  32. if "searchbot" not in f.lower() and "searchook" not in f.lower():
  33. continue
  34. out.extend(get_books_from_list(f))
  35. return out
  36. for f in files:
  37. if f.lower().endswith(".epub"):
  38. log.info("EPUB %s", f)
  39. new_fname = f.replace("epub", "mobi")
  40. p = subprocess.Popen(["ebook-convert", f, new_fname], stdout=subprocess.DEVNULL)
  41. # TODO log to file?
  42. p.wait()
  43. out.append(new_fname)
  44. out.append(f)
  45. # TODO make paths absolute?
  46. return out
  47. def handle_results(q, cb):
  48. while True:
  49. try:
  50. item = q.get(timeout=1)
  51. if item is None:
  52. break
  53. except queue.Empty:
  54. continue
  55. log.info("Got a result! %s", item)
  56. if item['type'] == 'status':
  57. cb((item['key'], item['status']))
  58. elif item['type'] == 'files':
  59. cb(handle_files(item['files']))
  60. class BookClient:
  61. workers = []
  62. def __init__(self, cb):
  63. self.q = queue.Queue()
  64. self.rq = queue.Queue()
  65. self.results_t = Thread(target=handle_results, args=(self.rq, cb))
  66. self.results_t.daemon = True
  67. self.results_t.start()
  68. self.create_worker()
  69. def create_worker(self):
  70. print("Spawning a new worker")
  71. worker = IRCClient(command_queue=self.q, results_queue=self.rq)
  72. worker.start()
  73. self.workers.append(worker)
  74. def free_workers(self):
  75. return any([not worker.busy for worker in self.workers])
  76. def request(self, mode, query):
  77. if not self.free_workers():
  78. self.create_worker()
  79. if mode not in [MODE_SEARCH, MODE_BOOK]:
  80. print('Invalid mode')
  81. return
  82. self.q.put({'query': query, 'mode': mode})
  83. print("Command acknowledged")
  84. def __del__(self):
  85. print("Cleaning up workers")
  86. for worker in self.workers:
  87. worker.stop()
  88. worker.join(timeout=5)
  89. print("Cleaning up results")
  90. self.rq.put(None)
  91. self.results_t.join(timeout=5)