Added suffix option.

Author: TCLamnidis

eigenstrat_database_tools.py

@@ -8,6 +8,7 @@
 parser._optionals.title = "Available options"
 parser.add_argument("-i", "--Input", type=str, metavar="<INPUT FILES PREFIX>", required=True, help="The desired input file prefix. Input files are assumed to be <INPUT PREFIX>.geno, <INPUT PREFIX>.snp and <INPUT PREFIX>.ind .")
 parser.add_argument("-o", "--Output", type=str, metavar="<OUTPUT FILES PREFIX>", required=False, help="The desired output file prefix. Three output files are created, <OUTPUT FILES PREFIX>.geno , <OUTPUT FILES PREFIX>.snp and <OUTPUT FILES PREFIX>.ind .")
+parser.add_argument("-s", "--Suffix", type = str, metavar = "<INPUT FILE SUFFIX>", required = False, default = '', help = "The suffix (if any) that follows .geno/.snp/.ind in the input files. For example, specifying '-s .txt' will treat <INPUT PREFIX>.{geno,snp,ind}.txt as the desired input files.")
 group = parser.add_mutually_exclusive_group(required=True)
 group2 = parser.add_mutually_exclusive_group(required=False)
 group.add_argument("-C", "--Check", type=argparse.FileType('r'), metavar="<INPUT FILE>", required=False, help="Check the -i .ind file and the second .ind file for duplicate individuals. Population assignment and/or individual sex are not checked, only individual names. Names are case sensitive.")
@@ -26,9 +27,9 @@
 	if args.Output is None:
 		parser.error("Output files (-o) must be specified when using -R/-E functions")
 
-IndFile = open(args.Input+".ind", "r")
-GenoFile = open(args.Input+".geno", "r")
-SnpFile = open(args.Input+".snp", "r")
+IndFile = open(args.Input+".ind"+args.Suffix, "r")
+GenoFile = open(args.Input+".geno"+args.Suffix, "r")
+SnpFile = open(args.Input+".snp"+args.Suffix, "r")
 
 
 #Check function
@@ -54,7 +55,7 @@
 #Check for errors in input files
 ##Check geno and snp compatibility
 lineNo = ""
-for line in sh.grep(sh.wc("-l", args.Input+".geno", args.Input+".snp"), args.Input):
+for line in sh.grep(sh.wc("-l", args.Input+".geno"+args.Suffix, args.Input+".snp"+args.Suffix), args.Input):
 	if lineNo=="":
 		lineNo=line.strip().split()[0]
 	elif lineNo==line.strip().split()[0]:
@@ -63,9 +64,9 @@
 		raise IOError("Input .snp and .geno files do not match.")
 
 ##Check geno and ind compatibility
-with open(args.Input+".geno", "r") as f:
+with open(args.Input+".geno"+args.Suffix, "r") as f:
 	for line in f:
-		if str(len(line.strip())) == sh.wc("-l", args.Input+".ind").strip().split()[0]:
+		if str(len(line.strip())) == sh.wc("-l", args.Input+".ind"+args.Suffix).strip().split()[0]:
 			break
 		else:
 			raise IOError("Input .ind and .geno files do not match.")
@@ -93,7 +94,7 @@
 Sex = {}
 Pop = {}
 for i in Samples:
-	for line in sh.grep(sh.cat("-n",args.Input+".ind"), "{}".format(i),_ok_code=[0,1]):
+	for line in sh.grep(sh.cat("-n",args.Input+".ind"+args.Suffix), "{}".format(i),_ok_code=[0,1]):
 		fields=line.strip().split()
 		if fields[1] == i:
 			Index[fields[1]]=(int(fields[0]) -1)
@@ -118,7 +119,7 @@
 	for i in Samples:
 		try: print (i, Sex[i], Pop[i], sep="\t", file=IndOutFile)
 		except KeyError:
-			print ("Individual \"",i,"\" not found in ",args.Input+".ind file.", sep="", file=sys.stderr)
+			print ("Individual \"",i,"\" not found in ",args.Input+".ind"+args.Suffix+" file.", sep="", file=sys.stderr)
 
 	for line in SnpFile:
 		print (line, end="", file =SnpOutFile)