forked from biopython/biopython
-
Notifications
You must be signed in to change notification settings - Fork 0
/
NewickIO.py
332 lines (272 loc) · 12.2 KB
/
NewickIO.py
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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
# Copyright (C) 2009 by Eric Talevich (eric.talevich@gmail.com)
# Based on Bio.Nexus, copyright 2005-2008 by Frank Kauff & Cymon J. Cox.
# All rights reserved.
# This code is part of the Biopython distribution and governed by its
# license. Please see the LICENSE file that should have been included
# as part of this package.
"""I/O function wrappers for the Newick file format.
See: http://evolution.genetics.washington.edu/phylip/newick_doc.html
"""
__docformat__ = "restructuredtext en"
import re
from Bio._py3k import StringIO
from Bio.Phylo import Newick
class NewickError(Exception):
"""Exception raised when Newick object construction cannot continue."""
pass
tokens = [
(r"\(", 'open parens'),
(r"\)", 'close parens'),
(r"[^\s\(\)\[\]\'\:\;\,]+", 'unquoted node label'),
(r"\:[0-9]*\.?[0-9]+([eE][+-]?[0-9]+)?", 'edge length'),
(r"\,", 'comma'),
(r"\[(\\.|[^\]])*\]", 'comment'),
(r"\'(\\.|[^\'])*\'", 'quoted node label'),
(r"\;", 'semicolon'),
(r"\n", 'newline'),
]
tokenizer = re.compile('(%s)' % '|'.join(token[0] for token in tokens))
token_dict = dict((name, re.compile(token)) for (token, name) in tokens)
# ---------------------------------------------------------
# Public API
def parse(handle, **kwargs):
"""Iterate over the trees in a Newick file handle.
:returns: generator of Bio.Phylo.Newick.Tree objects.
"""
return Parser(handle).parse(**kwargs)
def write(trees, handle, plain=False, **kwargs):
"""Write a trees in Newick format to the given file handle.
:returns: number of trees written.
"""
return Writer(trees).write(handle, plain=plain, **kwargs)
# ---------------------------------------------------------
# Input
def _parse_confidence(text):
if text.isdigit():
return int(text)
# NB: Could make this more consistent by treating as a percentage
# return int(text) / 100.
try:
return float(text)
# NB: This should be in [0.0, 1.0], but who knows what people will do
# assert 0 <= current_clade.confidence <= 1
except ValueError:
return None
def _format_comment(text):
return '[%s]' % (text.replace('[', '\\[').replace(']', '\\]'))
def _get_comment(clade):
if hasattr(clade, 'comment') and clade.comment:
return _format_comment(str(clade.comment))
else:
return ''
class Parser(object):
"""Parse a Newick tree given a file handle.
Based on the parser in `Bio.Nexus.Trees`.
"""
def __init__(self, handle):
self.handle = handle
@classmethod
def from_string(cls, treetext):
handle = StringIO(treetext)
return cls(handle)
def parse(self, values_are_confidence=False, comments_are_confidence=False, rooted=False):
"""Parse the text stream this object was initialized with."""
self.values_are_confidence = values_are_confidence
self.comments_are_confidence = comments_are_confidence
self.rooted = rooted
buf = ''
unicodeChecked = False
unicodeLines = ("\xef", "\xff", "\xfe", "\x00")
for line in self.handle:
if not unicodeChecked:
# check for unicode byte order marks on first line only,
# these lead to parsing errors (on Python 2)
if line.startswith(unicodeLines):
raise NewickError("The file or stream you attempted to parse includes "
"unicode byte order marks. You must convert it to "
"ASCII before it can be parsed.")
unicodeChecked = True
buf += line.rstrip()
if buf.endswith(';'):
yield self._parse_tree(buf)
buf = ''
if buf:
# Last tree is missing a terminal ';' character -- that's OK
yield self._parse_tree(buf)
def _parse_tree(self, text):
"""Parses the text representation into an Tree object."""
tokens = re.finditer(tokenizer, text.strip())
new_clade = self.new_clade
root_clade = new_clade()
current_clade = root_clade
entering_branch_length = False
lp_count = 0
rp_count = 0
for match in tokens:
token = match.group()
if token.startswith("'"):
# quoted label; add characters to clade name
current_clade.name = token[1:-1]
elif token.startswith('['):
# comment
current_clade.comment = token[1:-1]
if self.comments_are_confidence:
# Try to use this comment as a numeric support value
current_clade.confidence = _parse_confidence(current_clade.comment)
elif token == '(':
# start a new clade, which is a child of the current clade
current_clade = new_clade(current_clade)
entering_branch_length = False
lp_count += 1
elif token == ',':
# if the current clade is the root, then the external parentheses are missing
# and a new root should be created
if current_clade is root_clade:
root_clade = new_clade()
current_clade.parent = root_clade
# start a new child clade at the same level as the current clade
parent = self.process_clade(current_clade)
current_clade = new_clade(parent)
entering_branch_length = False
elif token == ')':
# done adding children for this parent clade
parent = self.process_clade(current_clade)
if not parent:
raise NewickError('Parenthesis mismatch.')
current_clade = parent
entering_branch_length = False
rp_count += 1
elif token == ';':
break
elif token.startswith(':'):
# branch length or confidence
value = float(token[1:])
if self.values_are_confidence:
current_clade.confidence = value
else:
current_clade.branch_length = value
elif token == '\n':
pass
else:
# unquoted node label
current_clade.name = token
if not lp_count == rp_count:
raise NewickError('Number of open/close parentheses do not match.')
# if ; token broke out of for loop, there should be no remaining tokens
try:
next_token = next(tokens)
raise NewickError('Text after semicolon in Newick tree: %s'
% next_token.group())
except StopIteration:
pass
self.process_clade(current_clade)
self.process_clade(root_clade)
return Newick.Tree(root=root_clade, rooted=self.rooted)
def new_clade(self, parent=None):
"""Returns a new Newick.Clade, optionally with a temporary reference
to its parent clade."""
clade = Newick.Clade()
if parent:
clade.parent = parent
return clade
def process_clade(self, clade):
"""Final processing of a parsed clade. Removes the node's parent and
returns it."""
if (clade.name and not (self.values_are_confidence or
self.comments_are_confidence)
and clade.confidence is None):
clade.confidence = _parse_confidence(clade.name)
if not clade.confidence is None:
clade.name = None
if hasattr(clade, 'parent'):
parent = clade.parent
parent.clades.append(clade)
del clade.parent
return parent
# ---------------------------------------------------------
# Output
class Writer(object):
"""Based on the writer in Bio.Nexus.Trees (str, to_string)."""
def __init__(self, trees):
self.trees = trees
def write(self, handle, **kwargs):
"""Write this instance's trees to a file handle."""
count = 0
for treestr in self.to_strings(**kwargs):
handle.write(treestr + '\n')
count += 1
return count
def to_strings(self, confidence_as_branch_length=False,
branch_length_only=False, plain=False,
plain_newick=True, ladderize=None, max_confidence=1.0,
format_confidence='%1.2f', format_branch_length='%1.5f'):
"""Return an iterable of PAUP-compatible tree lines."""
# If there's a conflict in the arguments, we override plain=True
if confidence_as_branch_length or branch_length_only:
plain = False
make_info_string = self._info_factory(plain,
confidence_as_branch_length, branch_length_only, max_confidence,
format_confidence, format_branch_length)
def newickize(clade):
"""Convert a node tree to a Newick tree string, recursively."""
label = clade.name or ''
if label:
unquoted_label = re.match(token_dict['unquoted node label'], label)
if (not unquoted_label) or (unquoted_label.end() < len(label)):
label = "'%s'" % label.replace('\\', '\\\\').replace("'", "\\'")
if clade.is_terminal(): # terminal
return (label
+ make_info_string(clade, terminal=True))
else:
subtrees = (newickize(sub) for sub in clade)
return '(%s)%s' % (','.join(subtrees),
label + make_info_string(clade))
# Convert each tree to a string
for tree in self.trees:
if ladderize in ('left', 'LEFT', 'right', 'RIGHT'):
# Nexus compatibility shim, kind of
tree.ladderize(reverse=(ladderize in ('right', 'RIGHT')))
rawtree = newickize(tree.root) + ';'
if plain_newick:
yield rawtree
continue
# Nexus-style (?) notation before the raw Newick tree
treeline = ['tree', (tree.name or 'a_tree'), '=']
if tree.weight != 1:
treeline.append('[&W%s]' % round(float(tree.weight), 3))
if tree.rooted:
treeline.append('[&R]')
treeline.append(rawtree)
yield ' '.join(treeline)
def _info_factory(self, plain, confidence_as_branch_length,
branch_length_only, max_confidence, format_confidence,
format_branch_length):
"""Return a function that creates a nicely formatted node tag."""
if plain:
# Plain tree only. That's easy.
def make_info_string(clade, terminal=False):
return _get_comment(clade)
elif confidence_as_branch_length:
# Support as branchlengths (eg. PAUP), ignore actual branchlengths
def make_info_string(clade, terminal=False):
if terminal:
# terminal branches have 100% support
return (':' + format_confidence % max_confidence) + _get_comment(clade)
else:
return (':' + format_confidence % clade.confidence) + _get_comment(clade)
elif branch_length_only:
# write only branchlengths, ignore support
def make_info_string(clade, terminal=False):
return (':' + format_branch_length % clade.branch_length) + _get_comment(clade)
else:
# write support and branchlengths (e.g. .con tree of mrbayes)
def make_info_string(clade, terminal=False):
if (terminal or
not hasattr(clade, 'confidence') or
clade.confidence is None):
return (':' + format_branch_length
) % (clade.branch_length or 0.0) + _get_comment(clade)
else:
return (format_confidence + ':' + format_branch_length
) % (clade.confidence, clade.branch_length or 0.0) + _get_comment(clade)
return make_info_string