bingo: elastic: new exact search, major refactorings

Author: AATDev21

bingo/bingo-elastic/python/bingo_elastic/elastic.py

@@ -22,7 +22,7 @@
 except ImportError:
     pass
 
-from indigo import Indigo  # type: ignore
+from indigo import Indigo, IndigoObject  # type: ignore
 
 from bingo_elastic.model.record import (
     IndigoRecord,
@@ -105,6 +105,8 @@ def get_client(
             "sub_fingerprint": {"type": "keyword", "similarity": "boolean"},
             "sub_fingerprint_len": {"type": "integer"},
             "cmf": {"type": "binary"},
+            "hash": {"type": "unsigned_long"},
+            "has_error": {"type": "integer"},
         }
     }
 }
@@ -151,13 +153,16 @@ def prepare(
 
 
 def response_to_records(
-    res, index_name, postprocess_actions
+    res: dict,
+    index_name: str,
+    postprocess_actions: PostprocessType = None,
+    indigo_session: Indigo = None,
+    options: str = "",
 ) -> Generator[IndigoRecord, None, None]:
-    indigo_session = Indigo()
     for el_response in res.get("hits", {}).get("hits", []):
         record = get_record_by_index(el_response, index_name)
-        for action_fn in postprocess_actions:
-            record = action_fn(record, indigo_session)  # type: ignore
+        for action_fn in postprocess_actions:  # type: ignore
+            record = action_fn(record, indigo_session, options)  # type: ignore
             if not record:
                 continue
         yield record
@@ -218,27 +223,25 @@ async def index_records(self, records: Generator, chunk_size: int = 500):
 
     async def filter(
         self,
-        similarity: Union[BaseMatch] = None,
-        exact: IndigoRecord = None,
-        substructure: IndigoRecord = None,
-        limit=10,
+        query_subject: Union[BaseMatch, IndigoObject, IndigoRecord] = None,
+        indigo_session: Indigo = None,
+        limit: int = 10,
+        options: str = "",
         **kwargs,
     ) -> AsyncGenerator[IndigoRecord, None]:
 
         # actions needed to be called on elastic_search result
         postprocess_actions: PostprocessType = []
 
         query = compile_query(
-            similarity=similarity,
-            exact=exact,
-            substructure=substructure,
+            query_subject=query_subject,
             limit=limit,
             postprocess_actions=postprocess_actions,
             **kwargs,
         )
         res = await self.el_client.search(index=self.index_name, body=query)
         for record in response_to_records(
-            res, self.index_name, postprocess_actions
+            res, self.index_name, postprocess_actions, indigo_session, options
         ):
             yield record
 
@@ -313,40 +316,33 @@ def delete_all_records(self):
 
     def filter(
         self,
-        similarity: Union[BaseMatch] = None,
-        exact: IndigoRecord = None,
-        substructure: IndigoRecord = None,
-        limit=10,
+        query_subject: Union[BaseMatch, IndigoObject, IndigoRecord] = None,
+        indigo_session: Indigo = None,
+        limit: int = 10,
+        options: str = "",
         **kwargs,
     ) -> Generator[IndigoRecord, None, None]:
 
         # actions needed to be called on elastic_search result
         postprocess_actions: PostprocessType = []
-
         query = compile_query(
-            similarity=similarity,
-            exact=exact,
-            substructure=substructure,
+            query_subject=query_subject,
             limit=limit,
             postprocess_actions=postprocess_actions,
             **kwargs,
         )
         res = self.el_client.search(index=self.index_name, body=query)
         yield from response_to_records(
-            res, self.index_name, postprocess_actions
+            res, self.index_name, postprocess_actions, indigo_session, options
         )
 
 
-# pylint: disable=too-many-arguments
 def compile_query(
-    similarity: BaseMatch = None,
-    exact: IndigoRecord = None,
-    substructure: IndigoRecord = None,
+    query_subject: Union[BaseMatch, IndigoObject, IndigoRecord] = None,
     limit: int = 10,
     postprocess_actions: PostprocessType = None,
     **kwargs,
 ) -> Dict:
-
     query = {
         "size": limit,
         "_source": {
@@ -359,17 +355,15 @@ def compile_query(
             ],
         },
     }
-    if similarity and substructure:
-        raise AttributeError(
-            "similarity and substructure search is not supported"
-        )
 
-    if similarity:
-        similarity.compile(query, postprocess_actions)
-    elif exact:
-        query_factory("exact", exact).compile(query, postprocess_actions)
-    elif substructure:
-        query_factory("substructure", substructure).compile(
+    if isinstance(query_subject, BaseMatch):
+        query_subject.compile(query, postprocess_actions)
+    elif isinstance(query_subject, IndigoRecord):
+        query_factory("exact", query_subject).compile(
+            query, postprocess_actions
+        )
+    elif isinstance(query_subject, IndigoObject):
+        query_factory("substructure", query_subject).compile(
             query, postprocess_actions
         )
 

bingo/bingo-elastic/python/bingo_elastic/model/record.py

@@ -5,6 +5,9 @@
 
 from indigo import Indigo, IndigoException, IndigoObject  # type: ignore
 
+MOL_TYPES = ["#02: <molecule>", "#03: <query reaction>", "#12: <RDFMolecule>"]
+REAC_TYPES = ["#04: <reaction>", "#05: <query reaction>"]
+
 
 # pylint: disable=unused-argument
 def skip_errors(instance: IndigoRecord, err: BaseException) -> None:
@@ -29,36 +32,60 @@ def __set__(self, instance: IndigoRecord, value: Dict):
 
 class WithIndigoObject:
     def __set__(self, instance: IndigoRecord, value: IndigoObject) -> None:
+        value_dup = value.clone()
+        value_dup.aromatize()
+
         fingerprints = (
             "sim",
             "sub",
         )
+
         for f_print in fingerprints:
             try:
                 setattr(instance, f"{f_print}_fingerprint", [])
                 setattr(instance, f"{f_print}_fingerprint_len", 0)
-                fp_ = [
-                    int(feature)
-                    for feature in value.fingerprint(f_print)
-                    .oneBitsList()
-                    .split(" ")
-                ]
-                setattr(instance, f"{f_print}_fingerprint", fp_)
-                setattr(instance, f"{f_print}_fingerprint_len", len(fp_))
+
+                fp_list = value_dup.fingerprint(f_print).oneBitsList()
+                if fp_list:
+                    fp_ = [int(feature) for feature in fp_list.split(" ")]
+                    setattr(instance, f"{f_print}_fingerprint", fp_)
+                    setattr(instance, f"{f_print}_fingerprint_len", len(fp_))
             except ValueError as err_:
                 check_error(instance, err_)
             except IndigoException as err_:
                 check_error(instance, err_)
 
         try:
-            setattr(instance, "name", value.name())
-        except IndigoException:
-            pass
+            cmf = " ".join(map(str, list(value_dup.serialize())))
+            setattr(instance, "cmf", cmf)
+        except IndigoException as err_:
+            setattr(instance, "cmf", "")
+            check_error(instance, err_)
+
+        try:
+            setattr(instance, "name", value_dup.name())
+        except IndigoException as err_:
+            setattr(instance, "name", "")
+            check_error(instance, err_)
+
+        try:
+            internal_type = value_dup.dbgInternalType()
+            if internal_type in MOL_TYPES:
+                hash_ = [
+                    component.clone().hash()
+                    for component in value_dup.iterateComponents()
+                ]
+                setattr(instance, "hash", sorted(set(hash_)))
+            elif internal_type in REAC_TYPES:
+                setattr(instance, "hash", [value_dup.hash()])
+        except IndigoException as err_:
+            check_error(instance, err_)
 
         try:
-            setattr(
-                instance, "cmf", " ".join(map(str, list(value.serialize())))
-            )
+            if value_dup.checkBadValence():
+                setattr(instance, "has_error", 1)
+            else:
+                setattr(instance, "has_error", 0)
         except IndigoException as err_:
             check_error(instance, err_)
 
@@ -74,6 +101,7 @@ class IndigoRecord:
 
     cmf: Optional[str] = None
     name: Optional[str] = None
+    rawData: Optional[str] = None
     sim_fingerprint: Optional[List[str]] = None
     sub_fingerprint: Optional[List[str]] = None
     indigo_object = WithIndigoObject()
@@ -120,8 +148,9 @@ def as_dict(self) -> Dict:
         }
 
     def as_indigo_object(self, session: Indigo):
-        assert self.cmf
-        return session.deserialize(list(map(int, self.cmf.split(" "))))
+        if self.cmf:
+            return session.deserialize(list(map(int, self.cmf.split(" "))))  # type: ignore
+        raise ValueError("Unexpected cmf value")
 
 
 class IndigoRecordMolecule(IndigoRecord):