[WIP] add zipfile collection support

src/sourmash/commands.py

@@ -519,6 +519,8 @@ def search(args):
 
 def categorize(args):
     "Use a database to find the best match to many signatures."
+    from .index import MultiIndex
+
     set_quiet(args.quiet)
     moltype = sourmash_args.calculate_moltype(args)
 
@@ -533,24 +535,27 @@ def categorize(args):
     # load search database
     tree = load_sbt_index(args.sbt_name)
 
-    # load query filenames
-    inp_files = set(sourmash_args.traverse_find_sigs(args.queries))
-    inp_files = inp_files - already_names
-
-    notify('found {} files to query', len(inp_files))
-
-    loader = sourmash_args.LoadSingleSignatures(inp_files,
-                                                args.ksize, moltype)
+    # utility function to load & select relevant signatures.
+    def _yield_all_sigs(queries, ksize, moltype):
+        for filename in queries:
+            mi = MultiIndex.load_from_path(filename, False)
+            mi = mi.select(ksize=ksize, moltype=moltype)
+            for ss, loc in mi.signatures_with_location():
+                yield ss, loc
 
     csv_w = None
     csv_fp = None
     if args.csv:
         csv_fp = open(args.csv, 'w', newline='')
         csv_w = csv.writer(csv_fp)
 
-    for queryfile, query, query_moltype, query_ksize in loader:
+    for query, loc in _yield_all_sigs(args.queries, args.ksize, moltype):
+        # skip if we've already done signatures from this file.
+        if loc in already_names:
+            continue
+
         notify('loaded query: {}... (k={}, {})', str(query)[:30],
-               query_ksize, query_moltype)
+               query.minhash.ksize, query.minhash.moltype)
 
         results = []
         search_fn = SearchMinHashesFindBest().search
@@ -575,14 +580,9 @@ def categorize(args):
             notify('for {}, no match found', query)
 
         if csv_w:
-            csv_w.writerow([queryfile, query, best_hit_query_name,
+            csv_w.writerow([loc, query, best_hit_query_name,
                            best_hit_sim])
 
-    if loader.skipped_ignore:
-        notify('skipped/ignore: {}', loader.skipped_ignore)
-    if loader.skipped_nosig:
-        notify('skipped/nosig: {}', loader.skipped_nosig)
-
     if csv_fp:
         csv_fp.close()
 

src/sourmash/index.py

@@ -1,10 +1,15 @@
 "An Abstract Base Class for collections of signatures."
 
+import sourmash
 from abc import abstractmethod, ABC
 from collections import namedtuple
+import zipfile
+import os
 
 
 class Index(ABC):
+    is_database = False
+
     @abstractmethod
     def signatures(self):
         "Return an iterator over all signatures in the Index object."
@@ -122,8 +127,55 @@ def gather(self, query, *args, **kwargs):
         return results
 
     @abstractmethod
-    def select(self, ksize=None, moltype=None):
-        ""
+    def select(self, ksize=None, moltype=None, scaled=None, num=None,
+               abund=None, containment=None):
+        """Return Index containing only signatures that match requirements.
+
+        Current arguments can be any or all of:
+        * ksize
+        * moltype
+        * scaled
+        * num
+        * containment
+
+        'select' will raise ValueError if the requirements are incompatible
+        with the Index subclass.
+
+        'select' may return an empty object or None if no matches can be
+        found.
+        """
+
+
+def select_signature(ss, ksize=None, moltype=None, scaled=0, num=0,
+                     containment=False):
+    "Check that the given signature matches the specificed requirements."
+    # ksize match?
+    if ksize and ksize != ss.minhash.ksize:
+        return False
+
+    # moltype match?
+    if moltype and moltype != ss.minhash.moltype:
+        return False
+
+    # containment requires scaled; similarity does not.
+    if containment:
+        if not scaled:
+            raise ValueError("'containment' requires 'scaled' in Index.select'")
+        if not ss.minhash.scaled:
+            return False
+
+    # 'scaled' and 'num' are incompatible
+    if scaled:
+        if ss.minhash.num:
+            return False
+    if num:
+        # note, here we check if 'num' is identical; this can be
+        # changed later.
+        if ss.minhash.scaled or num != ss.minhash.num:
+            return False
+
+    return True
+
 
 class LinearIndex(Index):
     "An Index for a collection of signatures. Can load from a .sig file."
@@ -155,23 +207,76 @@ def load(cls, location):
         lidx = LinearIndex(si, filename=location)
         return lidx
 
-    def select(self, ksize=None, moltype=None):
-        def select_sigs(ss, ksize=ksize, moltype=moltype):
-            if (ksize is None or ss.minhash.ksize == ksize) and \
-               (moltype is None or ss.minhash.moltype == moltype):
-               return True
+    def select(self, **kwargs):
+        """Return new LinearIndex containing only signatures that match req's.
 
-        return self.filter(select_sigs)
+        Does not raise ValueError, but may return an empty Index.
+        """
+        # eliminate things from kwargs with None or zero value
+        kw = { k : v for (k, v) in kwargs.items() if v }
 
-    def filter(self, filter_fn):
         siglist = []
         for ss in self._signatures:
-            if filter_fn(ss):
+            if select_signature(ss, **kwargs):
                 siglist.append(ss)
 
         return LinearIndex(siglist, self.filename)
 
 
+class ZipFileLinearIndex(Index):
+    """\
+    A read-only collection of signatures in a zip file.
+
+    Does not support `insert` or `save`.
+    """
+    is_database = True
+
+    def __init__(self, zf, selection_dict=None,
+                 traverse_yield_all=False):
+        self.zf = zf
+        self.selection_dict = selection_dict
+        self.traverse_yield_all = traverse_yield_all
+
+    @property
+    def filename(self):
+        return self.zf.filename
+
+    def insert(self, signature):
+        raise NotImplementedError
+
+    def save(self, path):
+        raise NotImplementedError
+
+    @classmethod
+    def load(cls, location, traverse_yield_all=False):
+        "Class method to load a zipfile."
+        zf = zipfile.ZipFile(location, 'r')
+        return cls(zf, traverse_yield_all=traverse_yield_all)
+
+    def signatures(self):
+        "Load all signatures in the zip file."
+        from .signature import load_signatures
+        for zipinfo in self.zf.infolist():
+            # should we load this file? if it ends in .sig OR we are forcing:
+            if zipinfo.filename.endswith('.sig') or self.traverse_yield_all:
+                fp = self.zf.open(zipinfo)
+
+                # now load all the signatures and select on ksize/moltype:
+                selection_dict = self.selection_dict
+                for ss in load_signatures(fp):
+                    if selection_dict:
+                        if select_signature(ss, **self.selection_dict):
+                            yield ss
+                    else:
+                        yield ss
+
+    def select(self, **kwargs):
+        "Select signatures in zip file based on ksize/moltype."
+        return ZipFileLinearIndex(self.zf,
+                                  selection_dict=kwargs,
+                                  traverse_yield_all=self.traverse_yield_all)
+
+
 class MultiIndex(Index):
     """An Index class that wraps other Index classes.
 
@@ -193,6 +298,11 @@ def signatures(self):
             for ss in idx.signatures():
                 yield ss
 
+    def signatures_with_location(self):
+        for idx, loc in zip(self.index_list, self.source_list):
+            for ss in idx.signatures():
+                yield ss, loc
+
     def __len__(self):
         return sum([ len(idx) for idx in self.index_list ])
 
@@ -203,14 +313,62 @@ def insert(self, *args):
     def load(self, *args):
         raise NotImplementedError
 
+    @classmethod
+    def load_from_path(cls, pathname, force=False):
+        "Create a MultiIndex from a path (filename or directory)."
+        from .sourmash_args import traverse_find_sigs
+        if not os.path.exists(pathname):
+            raise ValueError(f"'{pathname}' must be a directory")
+
+        index_list = []
+        source_list = []
+        for thisfile in traverse_find_sigs([pathname], yield_all_files=force):
+            try:
+                idx = LinearIndex.load(thisfile)
+                index_list.append(idx)
+                source_list.append(thisfile)
+            except (IOError, sourmash.exceptions.SourmashError):
+                if force:
+                    continue    # ignore error
+                else:
+                    raise       # contine past error!
+
+        db = None
+        if index_list:
+            db = cls(index_list, source_list)
+        else:
+            raise ValueError(f"no signatures to load under directory '{pathname}'")
+
+        return db
+
+    @classmethod
+    def load_from_file_list(cls, filename):
+        "Create a MultiIndex from all files listed in a text file."
+        from .sourmash_args import (load_file_list_of_signatures,
+                                    load_file_as_index)
+        idx_list = []
+        src_list = []
+
+        file_list = load_file_list_of_signatures(filename)
+        for fname in file_list:
+            idx = load_file_as_index(fname)
+            src = fname
+
+            idx_list.append(idx)
+            src_list.append(src)
+
+        db = MultiIndex(idx_list, src_list)
+        return db
+
     def save(self, *args):
         raise NotImplementedError
 
-    def select(self, ksize=None, moltype=None):
+    def select(self, **kwargs):
+        "Run 'select' on all indices within this MultiIndex."
         new_idx_list = []
         new_src_list = []
         for idx, src in zip(self.index_list, self.source_list):
-            idx = idx.select(ksize=ksize, moltype=moltype)
+            idx = idx.select(**kwargs)
             new_idx_list.append(idx)
             new_src_list.append(src)
 

src/sourmash/lca/command_summarize.py

@@ -10,6 +10,7 @@
 from ..logging import notify, error, print_results, set_quiet, debug
 from . import lca_utils
 from .lca_utils import check_files_exist
+from sourmash.index import MultiIndex
 
 
 DEFAULT_THRESHOLD=5
@@ -61,20 +62,15 @@ def load_singletons_and_count(filenames, ksize, scaled, ignore_abundance):
     total_count = 0
     n = 0
 
-    # in order to get the right reporting out of this function, we need
-    # to do our own traversal to expand the list of filenames, as opposed
-    # to using load_file_as_signatures(...)
-    filenames = sourmash_args.traverse_find_sigs(filenames)
-    filenames = list(filenames)
-
     total_n = len(filenames)
-
-    for query_filename in filenames:
+    for filename in filenames:
         n += 1
-        for query_sig in sourmash_args.load_file_as_signatures(query_filename,
-                                                               ksize=ksize):
+        mi = MultiIndex.load_from_path(filename)
+        mi = mi.select(ksize=ksize)
+
+        for query_sig, query_filename in mi.signatures_with_location():
             notify(u'\r\033[K', end=u'')
-            notify('... loading {} (file {} of {})', query_sig, n,
+            notify(f'... loading {query_sig} (file {n} of {total_n})',
                    total_n, end='\r')
             total_count += 1
 
@@ -87,7 +83,7 @@ def load_singletons_and_count(filenames, ksize, scaled, ignore_abundance):
             yield query_filename, query_sig, hashvals
 
     notify(u'\r\033[K', end=u'')
-    notify('loaded {} signatures from {} files total.', total_count, n)
+    notify(f'loaded {total_count} signatures from {n} files total.')
 
 
 def count_signature(sig, scaled, hashvals):