Repository: dlce-eva/python-newick

Issue:
Extended newick support
I'd like to process trees output by StarBEAST, but it uses nested annotations that aren't supported. An example 2 leaf tree:
```python

import newick

tree = "(1[&dmv={1},dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range={0.001,1.336},height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0,height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634},length_median=2.182,length_range={1.307,3.236}]:1.14538397925438,2[&dmv={1},dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range={0.001,1.336},height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0,height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634},length_median=2.182,length_range={1.307,3.236}]:1.14538397925438)[&dmv={1},dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range={0.001,1.336},height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0,height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634},length_median=2.182,length_range={1.307,3.236}]"

newick.loads(tree)
```
This fails with:

```
    newick.loads(tree)
  File "/home/jk/.local/lib/python3.8/site-packages/newick.py", line 369, in loads
    return [parse_node(ss.strip(), **kw) for ss in s.split(';') if ss.strip()]
  File "/home/jk/.local/lib/python3.8/site-packages/newick.py", line 369, in <listcomp>
    return [parse_node(ss.strip(), **kw) for ss in s.split(';') if ss.strip()]
  File "/home/jk/.local/lib/python3.8/site-packages/newick.py", line 471, in parse_node
    return Node.create(name=name, length=length, descendants=descendants, **kw)
  File "/home/jk/.local/lib/python3.8/site-packages/newick.py", line 76, in create
    node = cls(name=name, length=length, **kw)
  File "/home/jk/.local/lib/python3.8/site-packages/newick.py", line 42, in __init__
    raise ValueError(
ValueError: Node names or branch lengths must not contain ","
```

It should be possible to make the parser a bit more sophisticated using regex so that stuff nested in the comment brackets isn't split.

Would it be possible to add this feature, or is it outside the remit of the package? Thanks.

Comments:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 0000000..ed832ef
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,9 @@
+# Changes
+
+The `newick` package adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [v1.3.0] - 2021-05-04
+
+Added support for reading and writing of node annotations (in comments).
diff --git a/README.md b/README.md
index f9cd4c8..4f6f695 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,44 @@ a `list` of `newick.Node` objects.
 >>> trees = read(pathlib.Path('fname'))
 ```
 
+### Supported Newick dialects
+
+The ["Newick specification"](http://biowiki.org/wiki/index.php/Newick_Format) states
+
+> Comments are enclosed in square brackets and may appear anywhere
+
+This has spawned a host of ad-hoc mechanisms to insert additional data into Newick trees.
+
+The `newick` package allows to deal with comments in two ways.
+
+- Ignoring comments:
+  ```python
+  >>> newick.loads('[a comment](a,b)c;', strip_comments=True)[0].newick
+  '(a,b)c'
+  ```
+- Reading comments as node annotations: Several software packages use Newick comments to 
+  store node annotations, e.g. *BEAST, MrBayes or TreeAnnotator. Provided there are no
+  comments in places where they cannot be interpreted as node annotations, `newick` supports
+  reading and writing these annotations:
+  ```python
+  >>> newick.loads('(a[annotation],b)c;')[0].descendants[0].name
+  'a'
+  >>> newick.loads('(a[annotation],b)c;')[0].descendants[0].comment
+  'annotation'
+  >>> newick.loads('(a[annotation],b)c;')[0].newick
+  '(a[annotation],b)c'
+  ```
+
+Note that square brackets inside *quoted labels* will **not** be interpreted as comments
+or annotations:
+```python
+>>> newick.loads("('a[label]',b)c;")[0].descendants[0].name
+"'a[label]'"
+>>> newick.loads("('a[label]',b)c;")[0].newick
+"('a[label]',b)c"
+```
+
+
 ## Writing Newick
 
 In parallel to the read operations there are three functions to serialize a single `Node` object or a `list` of `Node`
diff --git a/src/newick.py b/src/newick.py
index ac31951..b27c958 100644
--- a/src/newick.py
+++ b/src/newick.py
@@ -28,7 +28,7 @@ class Node(object):
     descendants. It further has an ancestor, which is *None* if the node is the
     root node of a tree.
     """
-    def __init__(self, name=None, length=None, **kw):
+    def __init__(self, name=None, length=None, comment=None, **kw):
         """
         :param name: Node label.
         :param length: Branch length from the new node to its parent.
@@ -42,6 +42,7 @@ class Node(object):
                 raise ValueError(
                     'Node names or branch lengths must not contain "%s"' % char)
         self.name = name
+        self.comment = comment
         self._length = length
         self.descendants = []
         self.ancestor = None
@@ -63,7 +64,7 @@ class Node(object):
             self._length = self._length_formatter(length_)
 
     @classmethod
-    def create(cls, name=None, length=None, descendants=None, **kw):
+    def create(cls, name=None, length=None, descendants=None, comment=None, **kw):
         """
         Create a new `Node` object.
 
@@ -73,7 +74,7 @@ class Node(object):
         :param kw: Additonal keyword arguments are passed through to `Node.__init__`.
         :return: `Node` instance.
         """
-        node = cls(name=name, length=length, **kw)
+        node = cls(name=name, length=length, comment=comment, **kw)
         for descendant in descendants or []:
             node.add_descendant(descendant)
         return node
@@ -86,6 +87,8 @@ class Node(object):
     def newick(self):
         """The representation of the Node in Newick format."""
         label = self.name or ''
+        if self.comment:
+            label += '[{}]'.format(self.comment)
         if self._length:
             label += ':' + self._length
         descendants = ','.join([n.newick for n in self.descendants])
@@ -423,10 +426,16 @@ def write(tree, fname, encoding='utf8'):
 
 
 def _parse_name_and_length(s):
-    length = None
+    length, comment = None, None
     if ':' in s:
-        s, length = s.split(':', 1)
-    return s or None, length or None
+        parts = s.split(':')
+        if ']' not in parts[-1]:  # A ] in length doesn't make sense - the : must be in a comment.
+            s = ':'.join(parts[:-1])
+            length = parts[-1]
+    if '[' in s and s.endswith(']'):  # This looks like a node annotation in a comment.
+        s, comment = s.split('[', maxsplit=1)
+        comment = comment[:-1]
+    return s or None, length or None, comment
 
 
 def _parse_siblings(s, **kw):
@@ -434,11 +443,12 @@ def _parse_siblings(s, **kw):
     http://stackoverflow.com/a/26809037
     """
     bracket_level = 0
+    square_bracket_level = 0
     current = []
 
     # trick to remove special-case of trailing chars
     for c in (s + ","):
-        if c == "," and bracket_level == 0:
+        if c == "," and bracket_level == 0 and square_bracket_level == 0:
             yield parse_node("".join(current), **kw)
             current = []
         else:
@@ -446,6 +456,10 @@ def _parse_siblings(s, **kw):
                 bracket_level += 1
             elif c == ")":
                 bracket_level -= 1
+            elif c == "[":
+                square_bracket_level += 1
+            elif c == "]":
+                square_bracket_level -= 1
             current.append(c)
 
 
@@ -470,5 +484,5 @@ def parse_node(s, strip_comments=False, **kw):
             raise ValueError('unmatched braces %s' % parts[0][:100])
         descendants = list(_parse_siblings(')'.join(parts[:-1])[1:], **kw))
         label = parts[-1]
-    name, length = _parse_name_and_length(label)
-    return Node.create(name=name, length=length, descendants=descendants, **kw)
+    name, length, comment = _parse_name_and_length(label)
+    return Node.create(name=name, length=length, comment=comment, descendants=descendants, **kw)


Test patch:
diff --git a/tests/test_newick.py b/tests/test_newick.py
index 93dcace..24b4078 100644
--- a/tests/test_newick.py
+++ b/tests/test_newick.py
@@ -95,6 +95,7 @@ class TestNodeDescendantsFunctionality(unittest.TestCase):
 
 def test_read_write(tmp_path):
     trees = read(pathlib.Path(__file__).parent / 'fixtures' / 'tree-glottolog-newick.txt')
+    assert '[' in trees[0].descendants[0].name
     descs = [len(tree.descendants) for tree in trees]
     # The bookkeeping family has 391 languages
     assert descs[0] == 391
@@ -372,3 +373,22 @@ def test_prune_node():
     t2 = loads(tree)[0]
     t2.prune_by_names(["E"])
     assert t1.newick == t2.newick
+
+
+def test_with_comments():
+    nwk = "(1[x&dmv={1},dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range=" \
+          "{0.001,1.336},height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0," \
+          "height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634}," \
+          "length_median=2.182,length_range={1.307,3.236}]:1.14538397925438," \
+          "2[&dmv={1},dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range=" \
+          "{0.001,1.336},height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0," \
+          "height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634}," \
+          "length_median=2.182,length_range={1.307,3.236}]:1.14538397925438)[y&dmv={1}," \
+          "dmv1=0.260,dmv1_95%_hpd={0.003,0.625},dmv1_median=0.216,dmv1_range={0.001,1.336}," \
+          "height=1.310e-15,height_95%_hpd={0.0,3.552e-15},height_median=0.0," \
+          "height_range={0.0,7.105e-15},length=2.188,length_95%_hpd={1.725,2.634}," \
+          "length_median=2.182,length_range={1.307,3.236}]"
+    tree = loads(nwk)[0]
+    assert tree.comment.startswith('y')
+    assert tree.descendants[0].name == '1' and tree.descendants[0].comment.startswith('x')
+    assert tree.newick == nwk
\ No newline at end of file