Current Dev State

utilities/clients.py

@@ -0,0 +1,83 @@
+import json
+
+import requests
+
+
+class AMBITResult:
+
+    def __init__(self, *args, **kwargs):
+        self.smiles = kwargs['smiles']
+        self.tps = []
+        for bt in kwargs['products']:
+            if len(bt['products']):
+                self.tps.append(bt)
+
+        self.probs = None
+
+    def __str__(self):
+        x = self.smiles + "\n"
+        total_bts = len(self.tps)
+
+        for i, tp in enumerate(self.tps):
+            prob = ""
+            if self.probs:
+                prob = f" (p={self.probs[tp['id']]})"
+
+            if i == total_bts - 1:
+                x += f"\t└── {tp['name']}{prob}\n"
+            else:
+                x += f"\t├── {tp['name']}{prob}\n"
+
+            total_products = len(tp['products'])
+            for j, p in enumerate(tp['products']):
+                if j == total_products - 1:
+                    if i == total_bts - 1:
+                        x += f"\t\t└── {p}"
+                    else:
+                        x += f"\t│\t└── {p}\n"
+                else:
+                    if i == total_bts - 1:
+                        x += f"\t\t├── {p}\n"
+                    else:
+                        x += f"\t│\t├── {p}\n"
+        return x
+
+    def set_probs(self, probs):
+        self.probs = probs
+
+
+class AMBIT:
+
+    def __init__(self, host, rules=None):
+        self.host = host
+        self.rules = rules
+        self.ambit_params = {
+            'singlePos': True,
+            'split': False,
+        }
+
+    def batch_apply(self, smiles: list):
+        payload = {
+            'smiles': smiles,
+            'rules': self.rules,
+        }
+        payload.update(**self.ambit_params)
+
+        res = self._execute(payload)
+
+        tps = list()
+        for r in res['result']:
+            ar = AMBITResult(**r)
+            if len(ar.tps):
+                tps.append(ar)
+            else:
+                tps.append(None)
+        return tps
+
+    def apply(self, smiles: str):
+        return self.batch_apply([smiles])[0]
+
+    def _execute(self, payload):
+        res = requests.post(self.host + '/ambit', data=json.dumps(payload))
+        res.raise_for_status()
+        return res.json()