Read input with myopen and gzip

Author: enormandeau

fasta_extract.py

@@ -22,7 +22,7 @@ def write_to_file(self, handle):
         handle.write(self.sequence + "\n")
 
 # Functions
-def myopen(infile, mode="r"):
+def myopen(infile, mode="rt"):
     if infile.endswith(".gz"):
         return gzip.open(infile, mode=mode)
     else:
@@ -69,7 +69,7 @@ def fasta_iterator(input_file):
 
 fasta_sequences = fasta_iterator(fasta_file)
 
-with open(result_file, "w") as f:
+with open(result_file, "wt") as f:
     for seq in fasta_sequences:
         name = seq.name.split(" ")[0]
         if name in wanted and len(seq.sequence) > 0:

fasta_unwrap.py

@@ -6,6 +6,8 @@
     %program  input_file  output_file
 """
 
+# Modules
+import gzip
 import sys
 
 try:
@@ -14,13 +16,23 @@
     print("This program requires the Biopython library")
     sys.exit(0)
 
+# Defining functions
+def myopen(_file, mode="rt"):
+    if _file.endswith(".gz"):
+        return gzip.open(_file, mode=mode)
+
+    else:
+        return open(_file, mode=mode)
+
+# Parse user input
 try:
-    in_file = open(sys.argv[1], "rU")
-    out_file = open(sys.argv[2], "w")
+    in_file = myopen(sys.argv[1], "rt")
+    out_file = myopen(sys.argv[2], "wt")
 except:
     print(__doc__)
     sys.exit(0)
 
+# Treat sequences
 sequences = ([seq.id, seq.seq.tostring()] for seq in SeqIO.parse(in_file, 'fasta'))
 with open(sys.argv[2], "w") as out_file:
     for seq in sequences: