biotaphy_scripts/tree_reader.py at master · biotaphy/biotaphy_scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import string, sys
from node import Node
#from cnode import Node

"""
this takes a newick string as instr
and reads the string and makes the
nodes and returns the root node
"""
def read_tree_string(instr):
    root = None
    index = 0
    nextchar = instr[index]
    start = True
    keepgoing = True
    curnode = None
    while keepgoing == True:
        if nextchar == "(":
            if start == True:
                root = Node()
                curnode = root
                start = False
            else:
                newnode = Node()
                curnode.add_child(newnode)
                curnode = newnode
        elif nextchar == "[":
            note = ""
            index += 1
            nextchar = instr[index]
            while True:
                if nextchar == "]":
                    break
                index += 1
                note += nextchar
                nextchar = instr[index]
            curnode.note = note
        elif nextchar == ',':
            curnode = curnode.parent
        elif nextchar == ")":
            curnode = curnode.parent
            index += 1
            nextchar = instr[index]
            name = ""
            while True:
                if nextchar == ',' or nextchar == ')' or nextchar == ':' \
                    or nextchar == ';' or nextchar == '[':
                    break
                name += nextchar
                if index < len(instr)-1:
                    index += 1
                    nextchar = instr[index]
                else:
                    keepgoing = False
                    break
            curnode.label = name
            index -= 1
        elif nextchar == ';':
            keepgoing = False
            break
        elif nextchar == ":":
            index += 1
            nextchar = instr[index]
            brlen = ""
            while True:
                if nextchar == ',' or nextchar == ')' or nextchar == ':' \
                    or nextchar == ';' or nextchar == '[':
                    break
                brlen += nextchar
                index += 1
                nextchar = instr[index]
            curnode.length = float(brlen)
            index -= 1
        elif nextchar == ' ':
            index += 1
            nextchar = instr[index]
        else: # this is an external named node
            newnode = Node()
            curnode.add_child(newnode)
            curnode = newnode
            curnode.istip = True
            name = ""
            squote = False
            dquote = False
            if nextchar == "'":
                squote = True
                name += "'"
                index += 1
                nextchar = instr[index]
            elif nextchar == "\"":
                dquote = True
                name += "\""
                index += 1
                nextchar = instr[index]
            if squote == False and dquote == False:
                while True:
                    if nextchar == ',' or nextchar == ')' or nextchar == ':' \
                        or nextchar == ';' or nextchar == '[':
                        break
                    name += nextchar
                    index += 1
                    nextchar = instr[index]
            elif squote == True:
                while True:
                    if nextchar == "'":
                        name += "'"
                        index += 1
                        break
                    name += nextchar
                    index += 1
                    nextchar = instr[index]
            elif dquote == True:
                while True:
                    if nextchar == "\"":
                        name += "\""
                        index += 1
                        break
                    name += nextchar
                    index += 1
                    nextchar = instr[index]
            curnode.label = name
            index -= 1
        if index < len(instr) - 1:
            index += 1
        nextchar = instr[index]
    return root

def read_tree_file_iter(inf):
    info = open(inf,"r")
    for i in info:
        if len(i) > 2:
            yield read_tree_string(i.strip())
    info.close()


if __name__ == "__main__":
    s = "('a,daflkdasfljk\"dsa\"':3,(b:1e-05,c:1.3)int_|_and_33.5:5)root;"
    n2 = read_tree_string(s)
    print n2.get_newick_repr(True)