Merge pull request #27 from cov-lineages/dev

Dev

Author: rmcolq

scorpio/__init__.py

@@ -1,2 +1,2 @@
 _program = "scorpio"
-__version__ = "0.3.9"
+__version__ = "0.3.11"

scorpio/__main__.py

@@ -3,6 +3,7 @@
 import argparse
 import sys
 import os
+import logging
 
 import constellations
 
@@ -22,20 +23,18 @@ def main(sysargs = sys.argv[1:]):
     parser.add_argument("-cv", "--constellations-version", action='version', version=f"constellations {constellations.__version__}", help="show constellation's version number and exit")
 
     subparsers = parser.add_subparsers(
-        title="Available subcommands", help="", metavar=""
+        title="Available subcommands", help="", metavar="", dest='command'
     )
     # _______________________________   common  _________________________________#
     common = argparse.ArgumentParser(prog=_program, add_help=False)
 
     io_group = common.add_argument_group('Input/output options')
-    io_group.add_argument("-i", "--input", dest="input", required=True, help="Primary input file")
+    io_group.add_argument("-i", "--input", dest="input", required=False, help="Primary input file")
     io_group.add_argument("-m", "--metadata", dest="metadata", required=False, help="CSV of associated metadata")
-
     io_group.add_argument("-o", "--output", dest="output", required=False, help="Output file or path")
     io_group.add_argument("-p", "--prefix", dest="prefix", required=False, help="Output prefix. Default: scorpio")
     io_group.add_argument("--log-file", dest="log_file", metavar='<filename>', required=False,
                           help="Log file to use (otherwise uses stdout)")
-    io_group.add_argument("--config", action="store", help="Input config file", dest="config")
 
     constellation_group = common.add_argument_group('Constellation options')
     constellation_group.add_argument("-c", "--constellations", dest="constellations", required=False, nargs='+',
@@ -51,9 +50,6 @@ def main(sysargs = sys.argv[1:]):
                                      help="Extra mutations to type")
 
     misc_group = common.add_argument_group('Misc options')
-    misc_group.add_argument('--tempdir', action="store",
-                            help="Specify where you want the temporary stuff to go Default: $TMPDIR")
-    misc_group.add_argument("--no-temp", action="store_true", help="Output all intermediate files")
     misc_group.add_argument("--verbose", action="store_true", help="Print lots of stuff to screen")
     misc_group.add_argument("--dry-run", dest="dry_run", action="store_true", help="Quit after checking constellations and variants are AOK")
     misc_group.add_argument('-t', '--threads', action='store', dest="threads", type=int, help="Number of threads")
@@ -161,7 +157,20 @@ def main(sysargs = sys.argv[1:]):
 
     subparser_define.set_defaults(func=scorpio.subcommands.define.run)
 
+    # _______________________________  list  __________________________________#
+
+    subparser_list = subparsers.add_parser(
+        "list",
+        parents=[common],
+        help="Lists the constellations installed that would be typed/classified with the provided input options",
+    )
+    subparser_list.add_argument(
+        '--reference-json', dest="reference_json", help='JSON file containing keys "genome" with reference sequence '
+                                                        'and "proteins", "features" or "genes" with features of interest'
+                                                        ' and their coordinates'
+    )
 
+    subparser_list.set_defaults(func=scorpio.subcommands.list.run)
     # _________________________________________________________________________#
 
     args = parser.parse_args()
@@ -194,6 +203,20 @@ def main(sysargs = sys.argv[1:]):
     if not os.path.exists(args.prefix):
         os.mkdir(args.prefix)
 
+    ## format logging
+    format = '%(levelname)s: %(message)s'
+    if args.verbose:
+        level = logging.DEBUG
+    elif args.command == 'list':
+        level = logging.ERROR
+    else:
+        level = logging.INFO
+
+    if args.log_file:
+        logging.basicConfig(filename=args.log_file, level=level, format=format)
+    else:
+        logging.basicConfig(level=level, format=format)
+
     if not args.reference_json or not args.constellations:
         constellations_dir = constellations.__path__[0]
         reference_json = args.reference_json
@@ -202,7 +225,7 @@ def main(sysargs = sys.argv[1:]):
         constellation_subdirs = ["data", "definitions"]
         for dir in constellation_subdirs:
             data_dir = os.path.join(constellations_dir, dir)
-            print(f"Looking in {data_dir} for data files...")
+            logging.info(f"Looking in {data_dir} for data files...")
             for r, d, f in os.walk(data_dir):
                 for fn in f:
                     if fn == "SARS-CoV-2.json":
@@ -214,22 +237,21 @@ def main(sysargs = sys.argv[1:]):
                     elif not args.pangolin and fn.endswith(".csv"):
                         list_constellation_files.append(os.path.join(r, fn))
         if (not args.reference_json and reference_json == "") or (not args.constellations and list_constellation_files == []):
-            print(sfunk.cyan(
-                """Please either provide a reference JSON and constellation definition file, or check your environment 
-                to make sure that constellations has been properly installed."""))
+            logging.warning("""Please either provide a reference JSON and constellation definition file, or check your environment 
+                to make sure that constellations has been properly installed.""")
             sys.exit(-1)
         if not args.reference_json:
             args.reference_json = reference_json
-            print("Found reference %s" %args.reference_json)
+            logging.info("Found reference %s" %args.reference_json)
         if not args.constellations:
             args.constellations = list_constellation_files
-            print("Found constellations:")
+            logging.info("Found constellations:")
             for c in args.constellations:
-                print(c)
-            print("\n")
+                logging.info(c)
+            logging.info("\n")
 
-        if "call_all" in args and args.call_all and args.long:
-            print("Cannot provide long format summary file with multiple calls, ignoring --long\n")
+        if "call_all" in args and args.call_all and args.long and args.verbose:
+            logging.info("Cannot provide long format summary file with multiple calls, ignoring --long\n")
 
         if "append_genotypes" in args and args.append_genotypes and not args.ref_char:
             args.ref_char = None

scorpio/scripts/extract_definitions.py

@@ -5,6 +5,7 @@
 import operator
 import argparse
 import json
+import logging
 from Bio.Seq import Seq
 from operator import itemgetter
 
@@ -65,12 +66,12 @@ def get_group_dict(in_variants, group_column, index_column, subset):
                 if subset and row[group_column] not in subset:
                     continue
                 if row[index_column] in group_dict:
-                    print("%s is a duplicate in group CSV, keeping first")
+                    logging.warning("%s is a duplicate in group CSV, keeping first")
                 else:
                     group_dict[row[index_column]] = row[group_column]
                     groups.add(row[group_column])
 
-    print("Found", len(groups), "groups")
+    logging.info("Found %d groups" % len(groups))
 
     return group_dict
 
@@ -245,14 +246,14 @@ def extract_definitions(in_variants, in_groups, group_column, index_column, refe
     with open(in_variants, 'r', newline = '') as csv_in:
         reader = csv.DictReader(csv_in, delimiter=",", quotechar='\"', dialect = "unix")
         if index_column not in reader.fieldnames:
-            print("Index column %s not found in %s" % (index_column, in_variants))
+            logging.warning("Index column %s not found in %s" % (index_column, in_variants))
 
         if "nucleotide_variants" in reader.fieldnames:
             var_column = "nucleotide_variants"
         elif "nucleotide_mutations" in reader.fieldnames:
             var_column = "nucleotide_mutations"
         else:
-            print("No nucleotide_variants or nucleotide_mutations columns found")
+            logging.warning("No nucleotide_variants or nucleotide_mutations columns found")
             sys.exit(-1)
 
         for row in reader:
@@ -270,7 +271,7 @@ def extract_definitions(in_variants, in_groups, group_column, index_column, refe
                     for lineage in outgroup_dict[group_dict[index]]:
                         update_var_dict(outgroup_var_dict, lineage, variants)
             else:
-                print("Index column or variants column not in row", row)
+                logging.warning("Index column or variants column not in row", row)
 
     #print("outgroup_var_dict", outgroup_var_dict)
     #print("var_dict", var_dict)

scorpio/scripts/type_constellations.py

@@ -7,6 +7,7 @@
 import sys
 import json
 import re
+import logging
 
 if sys.version_info[0] < 3:
     raise Exception("Python 3 or a more recent version is required.")
@@ -54,7 +55,7 @@ def load_feature_coordinates(reference_json):
                         features_dict[name] = (start, end, json_dict[feature][item]["gene"])
                     else:
                         features_dict[name] = (start, end)
-                    print("Found reference feature %s with coordinates" % name, features_dict[name])
+                    logging.info("Found reference feature %s with coordinates" % name, features_dict[name])
     if len(features_dict) == 0:
         sys.stderr.write("No features (keys \"genes\", \"proteins\" or \"features\" ) provided in JSON %s " %
                          reference_json)
@@ -148,7 +149,7 @@ def variant_to_variant_record(l, refseq, features_dict):
         else:
             info["ref_start"] = info["pos"]
             info["name"] = l
-        print("Warning: found variant of type insertion, which will be ignored during typing")
+        logging.warning("Warning: found variant of type insertion, which will be ignored during typing")
     elif lsplit[0] in ["snp", "nuc"]:
         info = {"name": l, "type": "snp"}
         m = re.match(r'(?P<ref_allele>[ACGTUN]+)(?P<ref_start>\d+)(?P<alt_allele>[AGCTUN]*)', l[4:])
@@ -174,7 +175,7 @@ def variant_to_variant_record(l, refseq, features_dict):
         m = re.match(r'[aa:]*(?P<cds>\w+):(?P<ref_allele>[a-zA-Z-*]+)(?P<aa_pos>\d+)(?P<alt_allele>[a-zA-Z-*]*)', l)
         if not m:
             sys.stderr.write("Warning: couldn't parse the following string: %s - ignoring\n" % l)
-            # sys.exit(1)
+            sys.exit(1)
             return info
 
         info = m.groupdict()
@@ -255,7 +256,8 @@ def parse_json_in(refseq, features_dict, variants_file, constellation_names=None
     if constellation_names and name not in constellation_names:
         return variant_list, name, rules, mrca_lineage, incompatible_lineage_calls
 
-    print("\nParsing constellation JSON file %s" % variants_file)
+    logging.info("\n")
+    logging.info("Parsing constellation JSON file %s" % variants_file)
 
     if "sites" in json_dict:
         for site in json_dict["sites"]:
@@ -287,7 +289,8 @@ def parse_csv_in(refseq, features_dict, variants_file, constellation_names=None)
     if constellation_names and name not in constellation_names:
         return variant_list, name, compulsory
 
-    print("\nParsing constellation CSV file %s" % variants_file)
+    logging.info("\n")
+    logging.info("Parsing constellation CSV file %s" % variants_file)
 
     csv_in = open("%s" % variants_file, 'r')
     reader = csv.DictReader(csv_in, delimiter=",")
@@ -299,7 +302,7 @@ def parse_csv_in(refseq, features_dict, variants_file, constellation_names=None)
 
         if "id" not in reader.fieldnames:
             csv_in.close()
-            print("Warning: CSV headerline does not contain 'id': %s - ignoring" % reader.fieldnames)
+            logging.info("Warning: CSV headerline does not contain 'id': %s - ignoring" % reader.fieldnames)
             return variant_list, name, compulsory
 
     for row in reader:
@@ -331,8 +334,8 @@ def parse_textfile_in(refseq, features_dict, variants_file, constellation_names=
     name = parse_name_from_file(variants_file)
     if constellation_names and name not in constellation_names:
         return variant_list, name
-
-    print("\nParsing constellation text file %s" % variants_file)
+    logging.info("\n")
+    logging.info("Parsing constellation text file %s" % variants_file)
 
     with open("%s" % variants_file, "r") as f:
         for line in f:
@@ -380,7 +383,8 @@ def parse_variants_in(refseq, features_dict, variants_file, constellation_names=
 
 
 def parse_mutations_in(mutations_file):
-    print("\nParsing mutations file %s" % mutations_file)
+    logging.info("\n")
+    logging.info("Parsing mutations file %s" % mutations_file)
 
     mutations_list = []
     with open("%s" % mutations_file, "r") as f:
@@ -390,7 +394,7 @@ def parse_mutations_in(mutations_file):
                 if l.startswith('id'):
                     continue
                 mutations_list.append(l)
-    print("Found %d mutations" % len(mutations_list))
+    logging.info("Found %d mutations" % len(mutations_list))
     return mutations_list
 
 
@@ -472,7 +476,7 @@ def call_variant_from_fasta(record_seq, var, ins_char="?", oth_char=None, codon=
         if len(query) % 3 != 0:
             query = query_allele.replace("-","N")
         if len(query) % 3 != 0:
-            print("Warning: while typing variant %s (before,ref,after) = (%s,%s,%s) found sequence with query allele %s treated as %s. Handling by adding Ns which will result in ambiguous calls" %(var["name"], var["before"], var["ref_allele"], var["after"], query_allele, query))
+            logging.warning("Warning: while typing variant %s (before,ref,after) = (%s,%s,%s) found sequence with query allele %s treated as %s. Handling by adding Ns which will result in ambiguous calls" %(var["name"], var["before"], var["ref_allele"], var["after"], query_allele, query))
         query_allele = query
         while len(query_allele) % 3 != 0:
             query_allele += "N"
@@ -548,7 +552,7 @@ def counts_follow_rules(counts, rules):
             else:
                 counts["rules"] += 1
         else:
-            print("Warning: Ignoring rule %s:%s" % (rule, str(rules[rule])))
+            logging.warning("Warning: Ignoring rule %s:%s" % (rule, str(rules[rule])))
     return is_rule_follower
 
 def count_and_classify(record_seq, variant_list, rules):
@@ -607,10 +611,10 @@ def type_constellations(in_fasta, list_constellation_files, constellation_names,
             continue
         if len(variants) > 0:
             constellation_dict[constellation] = variants
-            print("Found file %s for constellation %s containing %i variants" % (
+            logging.info("Found file %s for constellation %s containing %i variants" % (
                     constellation_file, constellation, len([v["name"] for v in variants])))
         else:
-            print("Warning: %s is not a valid constellation file - ignoring" % constellation_file)
+            logging.warning("Warning: %s is not a valid constellation file - ignoring" % constellation_file)
     if mutations_list:
         new_mutations_list = []
         for entry in mutations_list:
@@ -619,7 +623,7 @@ def type_constellations(in_fasta, list_constellation_files, constellation_names,
             else:
                 new_mutations_list.append(entry)
         mutations_list = new_mutations_list
-        print("Typing provided mutations %s" % ",".join(mutations_list))
+        logging.info("Typing provided mutations %s" % ",".join(mutations_list))
         mutation_variants = parse_mutations(reference_seq, features_dict, mutations_list)
         if len(constellation_dict) == 1 and "mutations" not in constellation_dict:
             constellation = list(constellation_dict)[0]
@@ -697,19 +701,19 @@ def classify_constellations(in_fasta, list_constellation_files, constellation_na
         if constellation_names and constellation not in constellation_names:
             continue
         if not rules:
-            print("Warning: No rules provided to classify %s - ignoring" % constellation)
+            logging.warning("Warning: No rules provided to classify %s - ignoring" % constellation)
             continue
         else:
             rule_dict[constellation] = rules
         if len(variants) > 0:
             constellation_dict[constellation] = variants
-            print("Found file %s for constellation %s containing %i variants" % (
+            logging.info("Found file %s for constellation %s containing %i variants" % (
             constellation_file, constellation, len([v["name"] for v in variants])))
-            print("Rules", rule_dict[constellation])
+            logging.info("Rules %s" %rule_dict[constellation])
             mrca_lineage_dict[constellation] = mrca_lineage
             incompatible_dict[constellation] = incompatible_lineage_calls
         else:
-            print("Warning: %s is not a valid constellation file - ignoring" % constellation_file)
+            logging.warning("Warning: %s is not a valid constellation file - ignoring" % constellation_file)
 
     if mutations_list:
         new_mutations_list = []
@@ -803,6 +807,23 @@ def classify_constellations(in_fasta, list_constellation_files, constellation_na
             counts_out[constellation].close()
 
 
+def list_constellations(list_constellation_files, constellation_names, reference_json, label=None):
+
+    reference_seq, features_dict = load_feature_coordinates(reference_json)
+
+    list_of_constellations = []
+    for constellation_file in list_constellation_files:
+        constellation, variants, ignore, mrca_lineage, incompatible_lineage_calls = parse_variants_in(reference_seq, features_dict, constellation_file, constellation_names, label=label)
+        if not constellation:
+            continue
+        if constellation_names and constellation not in constellation_names:
+            continue
+        if len(variants) > 0 and mrca_lineage:
+            list_of_constellations.append(mrca_lineage)
+        elif len(variants) > 0:
+            list_of_constellations.append(constellation)
+    print("\n".join(list_of_constellations))
+
 def parse_args():
     parser = argparse.ArgumentParser(description="""Type an alignment at specific sites and classify with a barcode.""",
                                      formatter_class=argparse.RawTextHelpFormatter)

scorpio/subcommands/__init__.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
 
-__all__ = ["classify", "haplotype", "report", "define"]
+__all__ = ["classify", "haplotype", "report", "define", "list"]
 
 from scorpio.subcommands import *

scorpio/subcommands/list.py

@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+
+from scorpio.scripts.type_constellations import *
+
+
+def run(options):
+    list_constellations(options.constellations,
+                        options.names,
+                        options.reference_json,
+                        options.label)