Merge pull request #14 from martinghunt/check_res_genes

Check res genes

ariba/__init__.py

@@ -10,6 +10,7 @@
     'histogram',
     'link',
     'mapping',
+    'refcheck',
     'scaffold_graph',
     'tasks',
 ]

ariba/refcheck.py

@@ -0,0 +1,67 @@
+import os
+import pyfastaq
+
+class Error (Exception): pass
+
+
+class Checker:
+    def __init__(self, infile, min_length=1):
+        self.infile = os.path.abspath(infile)
+        if not os.path.exists(self.infile):
+            raise Error('File not found: "' + self.infile + '". Cannot continue')
+
+        self.min_length = min_length
+
+
+    def check(self, error_code_on_exit=None):
+        file_reader = pyfastaq.sequences.file_reader(self.infile)
+
+        for seq in file_reader:
+            if not seq.looks_like_gene():
+                return False, 'Not a gene', seq
+            elif len(seq) < self.min_length:
+                return False, 'Too short', seq
+
+        return True, None, None
+
+
+    def fix(self, outprefix):
+        file_reader = pyfastaq.sequences.file_reader(self.infile)
+        old2new_out = outprefix + '.rename'
+        fasta_out = outprefix + '.fa'
+        log_out = outprefix + '.log'
+        names = {}
+        old2new_out_fh = pyfastaq.utils.open_file_write(old2new_out)
+        fasta_out_fh = pyfastaq.utils.open_file_write(fasta_out)
+        log_out_fh = pyfastaq.utils.open_file_write(log_out)
+
+        for seq in file_reader:
+            if len(seq) < self.min_length:
+                print(seq.id, 'Too short. Skipping', sep='\t', file=log_out_fh)
+                continue
+
+            if not seq.looks_like_gene():
+                seq.revcomp()
+                if seq.looks_like_gene():
+                    print(seq.id, 'Reverse complemented', sep='\t', file=log_out_fh)
+                else:
+                    print(seq.id, 'Does not look like a gene. Skipping', sep='\t', file=log_out_fh)
+                    continue
+
+            original_id = seq.id
+            # replace unwanted characters with underscores
+            to_replace = ' '
+            seq.id = seq.id.translate(str.maketrans(to_replace, '_' *  len(to_replace)))
+
+            if seq.id in names:
+                names[seq.id] += 1
+                seq.id += '.' + str(names[seq.id])
+            else:
+                names[seq.id] = 1
+
+            print(original_id, seq.id, sep='\t', file=old2new_out_fh)
+            print(seq, file=fasta_out_fh)
+
+        pyfastaq.utils.close(fasta_out_fh)
+        pyfastaq.utils.close(log_out_fh)
+        pyfastaq.utils.close(old2new_out_fh)

ariba/tasks/refcheck.py

@@ -0,0 +1,23 @@
+import argparse
+import sys
+import ariba
+
+def run():
+    parser = argparse.ArgumentParser(
+        description = 'Check or fix resistance genes FASTA file',
+        usage = 'ariba refcheck [options] <infile>')
+    parser.add_argument('-m', '--min_length', type=int, help='Minimum length in nucleotides of gene [%(default)s]', metavar='INT', default=6)
+    parser.add_argument('-o', '--outprefix', help='Prefix of output files. If this option is used, a fixed file will be output, together with information on what was changed in the input file. If this option is not used, the script dies if any input sequence is not OK')
+    parser.add_argument('infile', help='Input file containing genes to be checked', metavar='Filename')
+    options = parser.parse_args()
+
+    checker = ariba.refcheck.Checker(options.infile, min_length=options.min_length)
+
+    if options.outprefix:
+        checker.fix(options.outprefix)
+    else:
+        ok, reason, seq = checker.check()
+        if not ok:
+            print('The following sequence not OK, for the reason:', reason)
+            print(seq)
+            sys.exit(1)

ariba/tests/data/refcheck_test_check_not_gene.fa

@@ -0,0 +1,2 @@
+>gene1
+TTGTGATGA

ariba/tests/data/refcheck_test_check_ok.fa

@@ -0,0 +1,2 @@
+>gene1
+TTGTGGTGA

ariba/tests/data/refcheck_test_check_too_short.fa

@@ -0,0 +1,2 @@
+>gene1
+TTGTGGTGA

ariba/tests/data/refcheck_test_fix_in.fa

@@ -0,0 +1,14 @@
+>gene1
+TTGTCGTAA
+>gene2
+TTGTCGTCGTCGTAA
+>gene3
+TTGTCGTCGTCGTCGTAA
+>gene3
+TTGTCGTCGTCGTCGTCGTAA
+>gene4 hello
+TTGTCGTCGTCGTCGTAA
+>revcomp
+TTACGACGACGACGACGACAA
+>not_a_gene
+TTGTAATAATAA

ariba/tests/data/refcheck_test_fix_out.fa

@@ -0,0 +1,10 @@
+>gene2
+TTGTCGTCGTCGTAA
+>gene3
+TTGTCGTCGTCGTCGTAA
+>gene3.2
+TTGTCGTCGTCGTCGTCGTAA
+>gene4_hello
+TTGTCGTCGTCGTCGTAA
+>revcomp
+TTGTCGTCGTCGTCGTCGTAA

ariba/tests/data/refcheck_test_fix_out.rename

@@ -0,0 +1,5 @@
+gene2	gene2
+gene3	gene3
+gene3	gene3.2
+gene4 hello	gene4_hello
+revcomp	revcomp

ariba/tests/refcheck_test.py

@@ -0,0 +1,45 @@
+import unittest
+import os
+import filecmp
+import pyfastaq
+from ariba import refcheck
+
+modules_dir = os.path.dirname(os.path.abspath(refcheck.__file__))
+data_dir = os.path.join(modules_dir, 'tests', 'data')
+
+
+class TestRefcheck(unittest.TestCase):
+    def test_check_pass(self):
+        '''test check file OK'''
+        infile = os.path.join(data_dir, 'refcheck_test_check_ok.fa')
+        c = refcheck.Checker(infile)
+        self.assertEqual(c.check(), (True, None, None))
+
+
+    def test_check_file_fail_not_gene(self):
+        '''test check file fail not a gene'''
+        infile = os.path.join(data_dir, 'refcheck_test_check_not_gene.fa')
+        c = refcheck.Checker(infile)
+        seq = pyfastaq.sequences.Fasta('gene1', 'TTGTGATGA')
+        self.assertEqual(c.check(), (False, 'Not a gene', seq))
+
+
+    def test_check_file_fail_too_short(self):
+        '''test check file fail short gene'''
+        infile = os.path.join(data_dir, 'refcheck_test_check_too_short.fa')
+        c = refcheck.Checker(infile, min_length=10)
+        seq = pyfastaq.sequences.Fasta('gene1', 'TTGTGGTGA')
+        self.assertEqual(c.check(), (False, 'Too short', seq))
+
+
+    def test_check_fix(self):
+        '''test fix'''
+        infile = os.path.join(data_dir, 'refcheck_test_fix_in.fa')
+        tmp_prefix = 'tmp.refcheck_test_fix.out'
+        c = refcheck.Checker(infile, min_length=10)
+        c.fix(tmp_prefix)
+        for x in ['fa', 'log', 'rename']:
+            expected = os.path.join(data_dir, 'refcheck_test_fix_out.' + x)
+            got = tmp_prefix + '.' + x
+            self.assertTrue(filecmp.cmp(expected, got, shallow=False))
+            os.unlink(got)

scripts/ariba

@@ -5,13 +5,15 @@ import sys
 
 
 tasks = {
+    'refcheck': 'Check or fix input genes FASTA',
     'run': 'Run the ARIBA local assembly pipeline',
     'flag': 'Translate the meaning of a flag output by the pipeline',
     'version': 'Print version and exit',
 }
 
 
 ordered_tasks = [
+    'refcheck',
     'run',
     'flag',
     'version',