--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/ete_gene_cnv.py	Thu Oct 31 07:48:59 2019 -0400
@@ -0,0 +1,75 @@
+from __future__ import print_function
+
+import argparse
+import collections
+
+from ete3 import PhyloTree
+
+
+def printTSV(myDict, colList=None):
+    """ Pretty print a list of dictionaries (myDict) as a dynamically sized table.
+    If column names (colList) aren't specified, they will show in random order.
+    Author: Thierry Husson - Use it as you want but don't blame me.
+    """
+    if not colList:
+        colList = list(myDict[0].keys() if myDict else [])
+
+    myList = [colList]
+
+    for item in myDict:
+        myList.append([str(item[col] if item[col] is not None else '') for col in colList])
+
+    for item in myList:
+        print(*item, sep="\t")
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Gene Copy Number Finder')
+    parser.add_argument('--genetree', required=True, help='GeneTree in nhx format')
+    parser.add_argument('--speciesorder', required=True, help='Comma-separated species list')
+    args = parser.parse_args()
+
+    species_list = args.speciesorder.split(",")
+    species_list = [_.strip() for _ in species_list]
+    table = []
+
+    with open(args.genetree, "r") as f:
+        # reads multiple gene tree line by line gene tree
+        for line in f:
+            # Remove empty NHX features that can be produced by TreeBest but break ete3
+            line = line.replace('[&&NHX]', '')
+
+            # reads single gene tree
+            genetree = PhyloTree(line)
+            leaves = genetree.get_leaf_names()
+
+            leaves_parts = [_.split("_") for _ in leaves]
+            for i, leaf_parts in enumerate(leaves_parts):
+                if len(leaf_parts) != 2:
+                    raise Exception("Leaf node '%s' is not in gene_species format" % leaves[i])
+
+            leaves_species = [_[1] for _ in leaves_parts]
+            species_counter = collections.Counter(leaves_species)
+
+            # Assign to ref_species the first element of species_list which
+            # appears in a leaf node
+            for ref_species in species_list:
+                if ref_species in species_counter:
+                    break
+            else:
+                raise Exception("None of the specified species was found in the GeneTree '%s'" % line)
+
+            # Find the gene of the (first) leaf node for the ref_species
+            for leaf_parts in leaves_parts:
+                if leaf_parts[1] == ref_species:
+                    species_counter['gene'] = leaf_parts[0]
+                    break
+
+            table.append(species_counter)
+
+    colList = ["gene"] + species_list
+    printTSV(table, colList)
+
+
+if __name__ == "__main__":
+    main()