[WIP] Refactor subject database and signature loading for search, gather, and multigather.

sourmash/_minhash.py

@@ -566,16 +566,19 @@ def is_molecule_type(self, molecule):
         """Check if this MinHash is a particular human-readable molecule type.
 
         Supports 'protein', 'dayhoff', 'hp', 'DNA'.
+        @CTB deprecate for 4.0?
         """
         if molecule.lower() not in ('protein', 'dayhoff', 'hp', 'dna'):
             raise ValueError("unknown moltype in query, '{}'".format(molecule))
-        if self.is_protein and molecule == 'protein':
-            return True
-        elif self.dayhoff and molecule == 'dayhoff':
-            return True
-        elif self.hp and molecule == 'hp':
-            return True
-        elif molecule.lower() == "dna" and self.is_dna:
-            return True
-
-        return False
+        return molecule == self.moltype
+
+    @property
+    def moltype(self):                    # TODO: test in minhash tests
+        if self.is_protein:
+            return 'protein'
+        elif self.dayhoff:
+            return 'dayhoff'
+        elif self.hp:
+            return 'hp'
+        else:
+            return 'DNA'

sourmash/commands.py

@@ -414,12 +414,37 @@ def search(args):
     from .search import search_databases
 
     set_quiet(args.quiet)
-    moltype = sourmash_args.calculate_moltype(args)
 
-    # set up the query.
-    query = sourmash_args.load_query_signature(args.query,
-                                               ksize=args.ksize,
-                                               select_moltype=moltype)
+    # load the command line argument selectors & query signature(s).
+    db_loader = sourmash_args.SearchDBLoader2(require_scaled=True)
+    db_loader.parse_args_selectors(args)
+    if not db_loader.load_query(args.query):
+        error("Cannot load query signature; exiting.")
+        sys.exit(-1)
+    if not db_loader.check_query_against_arg_selectors():
+        error("Cannot match query signature with args; exiting.")
+        sys.exit(-1)
+
+    # load the databases.
+    loaded_db_list = []
+    for filename in args.databases:
+        db, params = sourmash_args.load_target_with_params(filename)
+        if not db:
+            error("couldn't load {}", filename)
+            sys.exit(-1)
+        if not db_loader.add_database(filename, params):
+            error("couldn't match {}", filename)
+            sys.exit(-1)
+
+        loaded_db_list.append((db, filename, 'XYZ'))
+
+    # now that we've loaded the databases, figure out which query (queries?)
+    # are compatible. If there's exactly one, perfect!
+    if not db_loader.decide_query():
+        error("couldn't find acceptable query.")
+        sys.exit(-1)
+
+    query = db_loader.chosen_query
     notify('loaded query: {}... (k={}, {})', query.name()[:30],
                                              query.minhash.ksize,
                                              sourmash_args.get_moltype(query))
@@ -435,15 +460,19 @@ def search(args):
                    query.minhash.scaled, int(args.scaled))
         query.minhash = query.minhash.downsample_scaled(args.scaled)
 
-    # set up the search databases
-    databases = sourmash_args.load_dbs_and_sigs(args.databases, query,
-                                                not args.containment,
-                                                args.traverse_directory)
-
     # forcibly ignore abundances if query has no abundances
     if not query.minhash.track_abundance:
         args.ignore_abundance = True
 
+    # set up the search databases
+    # now that we have the query, apply the same selector to the databaess.
+    databases = []
+    for (db, filename, _) in loaded_db_list:
+        new_db = db.select(ksize=query.minhash.ksize,
+                           moltype=query.minhash.moltype)
+        # @CTB: here is also where we select the scaled.
+        databases.append((new_db, filename, 'XXX'))
+
     if not len(databases):
         error('Nothing found to search!')
         sys.exit(-1)
@@ -570,17 +599,42 @@ def gather(args):
     from .search import gather_databases, format_bp
 
     set_quiet(args.quiet, args.debug)
-    moltype = sourmash_args.calculate_moltype(args)
 
-    # load the query signature & figure out all the things
-    query = sourmash_args.load_query_signature(args.query,
-                                               ksize=args.ksize,
-                                               select_moltype=moltype)
+    # load the command line argument selectors & query signature(s).
+    db_loader = sourmash_args.SearchDBLoader2(require_scaled=True)
+    db_loader.parse_args_selectors(args)
+    if not db_loader.load_query(args.query):
+        error("Cannot load query signature; exiting.")
+        sys.exit(-1)
+    if not db_loader.check_query_against_arg_selectors():
+        error("Cannot match query signature with args; exiting.")
+        sys.exit(-1)
+
+    # load the databases.
+    loaded_db_list = []
+    for filename in args.databases:
+        db, params = sourmash_args.load_target_with_params(filename)
+        if not db:
+            error("couldn't load {}", filename)
+            sys.exit(-1)
+        if not db_loader.add_database(filename, params):
+            error("couldn't match {}", filename)
+            sys.exit(-1)
+
+        loaded_db_list.append((db, filename, 'XYZ'))
+
+    # now that we've loaded the databases, figure out which query (queries?)
+    # are compatible. If there's exactly one, perfect!
+    if not db_loader.decide_query():
+        error("couldn't find acceptable query.")
+        sys.exit(-1)
+
+    query = db_loader.chosen_query
     notify('loaded query: {}... (k={}, {})', query.name()[:30],
                                              query.minhash.ksize,
                                              sourmash_args.get_moltype(query))
 
-    # verify signature was computed right.
+    # verify signature was computed with --scaled.
     if query.minhash.scaled == 0:
         error('query signature needs to be created with --scaled')
         sys.exit(-1)
@@ -596,14 +650,20 @@ def gather(args):
         error('no query hashes!? exiting.')
         sys.exit(-1)
 
-    # set up the search databases
-    databases = sourmash_args.load_dbs_and_sigs(args.databases, query, False,
-                                                 args.traverse_directory)
+    # now that we have the query, apply the same selector to the databaess.
+    databases = []
+    for (db, filename, _) in loaded_db_list:
+        new_db = db.select(ksize=query.minhash.ksize,
+                           moltype=query.minhash.moltype)
+        # @CTB: here is also where we select the scaled.
+        databases.append((new_db, filename, 'XXX'))
 
     if not len(databases):
         error('Nothing found to search!')
         sys.exit(-1)
 
+    ### execute the gather algorithm.
+
     found = []
     weighted_missed = 1
     new_max_hash = query.minhash.max_hash

sourmash/index.py

@@ -96,6 +96,9 @@ def gather(self, query, *args, **kwargs):
 
         return results
 
+    @abstractmethod
+    def select(self, ksize=None, moltype=None):
+        ""
 
 class LinearIndex(Index):
     def __init__(self, _signatures=None, filename=None):
@@ -125,3 +128,15 @@ def load(cls, location):
 
         lidx = LinearIndex(si, filename=location)
         return lidx
+
+    def select(self, ksize=None, moltype=None):
+        def select_sigs(siglist, ksize, moltype):
+            for ss in siglist:
+#                print(self.filename, ss, ss.minhash.ksize, ss.minhash.moltype,
+#                      ksize, moltype)
+                if (ksize is None or ss.minhash.ksize == ksize) and \
+                   (moltype is None or ss.minhash.moltype == moltype):
+                   yield ss
+
+        siglist=select_sigs(self._signatures, ksize, moltype)
+        return LinearIndex(siglist, self.filename)

sourmash/lca/lca_utils.py

@@ -199,6 +199,18 @@ def signatures(self):
         for v in self._signatures.values():
             yield SourmashSignature(v)
 
+    def select(self, ksize=None, moltype=None):
+        ok = True
+        if ksize is not None and self.ksize != ksize:
+            ok = False
+        if moltype is not None and moltype != 'DNA':
+            ok = False
+
+        if ok:
+            return self
+
+        raise ValueError("cannot select LCA on ksize {} / moltype {}".format(ksize, moltype))
+
     def load(self, db_name):
         "Load from a JSON file."
         xopen = open

sourmash/sbt.py

@@ -137,6 +137,20 @@ def signatures(self):
         for k in self.leaves():
             yield k.data
 
+    def select(self, ksize=None, moltype=None):
+        first_sig = next(iter(self.signatures()))
+
+        ok = True
+        if ksize is not None and first_sig.minhash.ksize != ksize:
+            ok = False
+        if moltype is not None and first_sig.minhash.moltype != moltype:
+            ok = False
+
+        if ok:
+            return self
+
+        raise ValueError("cannot select SBT on ksize {} / moltype {}".format(ksize, moltype))
+
     def new_node_pos(self, node):
         if not self._nodes:
             self.next_node = 1

sourmash/search.py

@@ -81,6 +81,7 @@ def _find_best(dblist, query, threshold_bp):
     best_filename = None
 
     # search across all databases
+    # @CTB filetype no longer needed here
     for (obj, filename, filetype) in dblist:
         for cont, match, fname in obj.gather(query, threshold_bp=threshold_bp):
             assert cont                   # all matches should be nonzero.