added pickle database handling and added better replies for some scenarios

1 files changed, 74 insertions, 38 deletions

diff --git a/plugin.py b/plugin.py
index 2677f8c..9c50c09 100644
--- a/plugin.py
+++ b/plugin.py
@@ -28,14 +28,15 @@
 
 ###
 
-from supybot import utils, plugins, ircutils, callbacks
+from supybot import utils, plugins, ircutils, callbacks, world, conf, log
 from supybot.commands import *
-from pathlib import Path
-from os import path
 import dateutil.parser
 import json
 import requests
+import pickle
+import sys
 from datetime import datetime
+import time
 try:
     from supybot.i18n import PluginInternationalization
     _ = PluginInternationalization('Tripsit')
@@ -43,17 +44,14 @@ except ImportError:
     # Placeholder that allows to run the plugin on a bot
     # without the i18n module
     _ = lambda x: x
-dose_filename = 'dose.json'
-dose_db = Path(dose_filename)
-if not path.isfile(dose_filename):
-    dose_db.write_text('{}')
-dose_data = json.loads(dose_db.read_text())
 
-URL_DRUG = "http://tripbot.tripsit.me/api/tripsit/getDrug"
-URL_COMBO = "http://tripbot.tripsit.me/api/tripsit/getInteraction"
-URL_WIKI = "http://drugs.tripsit.me/%s"
+filename = conf.supybot.directories.data.dirize("Tripsit.db")
 
-INSUFFLATED = ["Insufflated", "Insufflated-IR", "Insufflated-XR"]
+
+url_drug = "http://tripbot.tripsit.me/api/tripsit/getDrug"
+url_combo = "http://tripbot.tripsit.me/api/tripsit/getInteraction"
+
+insufflated = ["Insufflated", "Insufflated-IR", "Insufflated-XR"]
 
 METHODS = {
     "iv": ["IV"],
@@ -63,8 +61,8 @@ METHODS = {
 
     "oral": ["Oral", "Oral-IR", "Oral-XR"],
 
-    "insufflated": INSUFFLATED,
-    "snorted": INSUFFLATED,
+    "insufflated": insufflated,
+    "snorted": insufflated,
 
     "smoked": ["Smoked"]
 }
@@ -72,54 +70,88 @@ METHODS = {
 class Tripsit(callbacks.Plugin):
     """Harm-Reduction tools from tripsit's tripbot and the tripsitwiki"""
     threaded = True
+
+    def __init__(self, irc):
+        self.__parent = super(Tripsit, self)
+        self.__parent.__init__(irc)
+        self.db = {}
+        self._loadDb()
+        world.flushers.append(self._flushDb)
+
+    def _loadDb(self):
+        """Loads the (flatfile) database mapping nicks to doses."""
+
+        try:
+            with open(filename, "rb") as f:
+                self.db = pickle.load(f)
+        except Exception as e:
+            self.log.debug("Tripsit: Unable to load pickled database: %s", e)
+
+    def _flushDb(self):
+        """Flushes the (flatfile) database mapping nicks to doses."""
+
+        try:
+            with open(filename, "wb") as f:
+                pickle.dump(self.db, f, 2)
+        except Exception as e:
+            self.log.warning("Tripsit: Unable to write pickled database: %s", e)
+
+    def die(self):
+        self._flushDb()
+        world.flushers.remove(self._flushDb)
+        self.__parent.die()
+
     @wrap(['something', optional('something')])
     def drug(self, irc, msg, args, name, category):
         """<drug> [<category>]
         fetches data on drug from tripsit wiki
         """
         category_list = []
-        r = requests.get(URL_DRUG, params={f"name": name}).json()
-        # r = json.loads(utils.web.getUrlContent(URL_DRUG, data={f"name": name}))
+        r = requests.get(url_drug, params={"name": name}).json()
         if not r['err']:
             drug = r["data"][0]["pretty_name"]
             properties = r["data"][0]["properties"]
+            for key in properties:
+                category_list.append(key)
             if category == None:
-                # category_list = list(properties)
-                for key in properties:
-                    category_list.append(key)
-                re = ", ".join(category_list)
-                # re = category_list
+                re = drug + " Available categories are: " + ", ".join(category_list)
                 irc.reply(re)
             else:
                 if category in properties.keys():
-                    re = properties[category]
+                    re = drug + " " + properties[category]
                     irc.reply(re)
                 else:
-                    irc.error("Unknown category")
+                    irc.error(f"Unknown category {drug} Available categories are: " + ", ".join(category_list))
         else:
             irc.error("unknown drug")
 
     def combo(self, irc, msg, args, drugA, drugB):
         """<drugA> <drugB>
-        fetches known interaction between the substances provided.
+        fetches known interactions between the substances provided.
         """
-        r = requests.get(URL_COMBO, params={f"drugA": drugA, f"drugB": drugB}).json()
+        r = requests.get(url_combo, params={f"drugA": drugA, f"drugB": drugB}).json()
         if not r["err"] and r["data"][0]:
             interaction = r["data"][0]
             drug_a = interaction["interactionCategoryA"]
             drug_b = interaction["interactionCategoryB"]
             interaction_status = interaction["status"]
-            note = interaction["note"]
-            re = f"{drug_a} and {drug_b}: {interaction_status}, Note: {note}"
-            irc.reply(re)
+            re = f"{drug_a} and {drug_b}: {interaction_status}"
+            if 'note' in interaction:
+                note = interaction["note"]
+                re += f'. Note: {note}'
+                irc.reply(re)
+            else:
+                irc.reply(re)
         else:
-            irc.error("Unknown Combo")
+            irc.reply("Unknown combo (that doesn't mean it's safe). Known combos: lsd, mushrooms, dmt, mescaline, dox, nbomes, 2c-x, 2c-t-x, amt, 5-meo-xxt, cannabis, ketamine, mxe, dxm, pcp, nitrous, amphetamines, mdma, cocaine, caffeine, alcohol, ghb/gbl, opioids, tramadol, benzodiazepines, maois, ssris.")
+
     combo = wrap(combo, [("something"), ("something")])
 
     def idose(self, irc, msg, args, dose, name, method):
-        """<amt> <drug> <method> logs a dose to keep track of
+        """<amountt> <drug> <method>
+        logs a dose in the database for retrieval
         """
-        r = requests.get(URL_DRUG, params={f"name": name}).json()
+        r = requests.get(url_drug, params={"name": name}).json()
         found_method = False
         onset = None
         if not r['err']:
@@ -152,8 +184,8 @@ class Tripsit(callbacks.Plugin):
             drug_and_method = "%s via %s" % (drug_and_method, method)
 
         time = datetime.utcnow()
-        dose_data[str(msg.nick)] = { 'time': str(time), 'dose': dose, 'drug': drug_name, 'method': method }
-        re = f"Dosed {dose} of {drug_and_method} at {str(time)}"
+        self.db[msg.nick] = { 'time': str(time), 'dose': dose, 'drug': drug_name, 'method': method }
+        re = f"{msg.nick} dosed {dose} of {drug_and_method} at {str(time)}"
 
         if not onset == None:
             re += f". You should start feeling effects {onset} from now"
@@ -163,13 +195,17 @@ class Tripsit(callbacks.Plugin):
     def lastdose(self, irc, msg, args):
         """ retrieves saved dose
         """
-        if str(msg.nick) in dose_data:
-            lastdose = dose_data[str(msg.nick)]
-            time = dateutil.parser.isoparse(lastdose['time'])
-            re = f"You last dosed {lastdose['dose']} of {lastdose['drug']} via {lastdose['method']} at {time}"
+        if msg.nick in self.db:
+            lastdose = self.db[msg.nick]
+            time = datetime.utcnow()
+            dose_time = dateutil.parser.isoparse(lastdose['time'])
+            since_dose = time - dose_time
+            since_dose_seconds = since_dose.total_seconds()
+            since_dose_formatted = utils.str.format('%T', since_dose_seconds)
+            re = f"You last dosed {lastdose['dose']} of {lastdose['drug']} via {lastdose['method']} at {dose_time} {since_dose_formatted} ago"
             irc.reply(re)
         else:
-            irc.error('No last dose saved for you')
+            irc.error(f'No last dose saved for {msg.nick}')
 
     lastdose = wrap(lastdose)