taxonomy.named function can return all is_valid taxids

fhcrc · Nov 17, 2023 · 24727db · 24727db
1 parent bc1bdeb
commit 24727db
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 21 deletions.
diff --git a/taxtastic/subcommands/named.py b/taxtastic/subcommands/named.py
@@ -57,23 +57,18 @@ def build_parser(parser):
 def action(args):
     engine = sqlalchemy.create_engine(args.url, echo=args.verbosity > 3)
     tax = Taxonomy(engine, schema=args.schema)
-    if args.tax_ids:
-        tax_ids = args.tax_ids
-    elif args.tax_id_file:
-        tax_ids = (i.strip() for i in args.tax_id_file)
-        tax_ids = [i for i in tax_ids if i]
-    elif args.seq_info:
-        seq_info = csv.DictReader(args.seq_info)
-        tax_ids = (row['tax_id'] for row in seq_info)
-    named = set(tax.named(set(tax_ids), no_rank=not args.ranked))
     if args.seq_info:
+        named = set(tax.named(no_rank=not args.ranked))
+        seq_info = csv.DictReader(args.seq_info)
         out = csv.DictWriter(args.outfile, fieldnames=seq_info.fieldnames)
         out.writeheader()
-        args.seq_info.seek(0)
-        for i in csv.DictReader(args.seq_info):
-            if i['tax_id'] in named:
-                out.writerow(i)
+        out.writerows(i for i in seq_info if i['tax_id'] in named)
     else:
-        for i in tax_ids:
-            if i in named:
-                args.outfile.write(i + '\n')
+        if args.tax_ids:
+            tax_ids = args.tax_ids
+        else:
+            tax_ids = (i.strip() for i in args.tax_id_file)
+            tax_ids = [i for i in tax_ids if i]
+        named = set(tax.named(set(tax_ids), no_rank=not args.ranked))
+        tax_ids = (i for i in tax_ids if i in named)
+        args.outfile.write('\n'.join(tax_ids) + '\n')
diff --git a/taxtastic/taxonomy.py b/taxtastic/taxonomy.py
@@ -825,12 +825,11 @@ def species_below(self, tax_id):
             assert self.is_ancestor_of(newc, tax_id)
             return newc
 
-    def named(self, taxids, no_rank=True):
+    def named(self, taxids=None, no_rank=True):
         nodes = self.nodes
-        s = select(nodes.c.tax_id).\
-            where(and_(
-                nodes.c.tax_id.in_(taxids),
-                nodes.c.is_valid))
+        s = select(nodes.c.tax_id).where(nodes.c.is_valid)
+        if taxids:
+            s = s.where(nodes.c.tax_id.in_(taxids))
         if not no_rank:
             s = s.where(nodes.c.rank != 'no_rank')
         return [f[0] for f in self.fetchall(s)]