support writing top-level comments

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changes
 
+## [Unreleased]
+
+- Fixed support for writing comments.
+
+
 ## [v1.9.0] - 2023-11-17
 
 - Support renaming taxa when normalising NEXUS.

docs/nexus.rst

@@ -59,3 +59,23 @@ obtained by calling the generic factory method
 or specific implementations of `Block.from_data`, such as
 :meth:`commonnexus.blocks.characters.Characters.from_data` or
 :meth:`commonnexus.blocks.trees.Trees.from_data`
+
+
+Comments
+~~~~~~~~
+
+The ``from_data`` methods of blocks accept a keyword argument ``comment`` to add a comment to a
+block construct.
+
+To add a comment to the top of a NEXUS file, one can proceed as follows:
+
+.. code-block:: python
+
+    >>> nex = Nexus('#NEXUS\n[{}]\n'.format('Comment goes here'))
+    >>> nex.append_block(Block.from_commands([], name='theblock'))
+    >>> print(nex)
+    #NEXUS
+    [Comment goes here]
+
+    BEGIN theblock;
+    END;

src/commonnexus/nexus.py

@@ -5,6 +5,7 @@
 import dataclasses
 
 from .tokenizer import TokenType, iter_tokens, get_name
+from .util import log_or_raise
 from commonnexus.command import Command
 from commonnexus.blocks import Block
 
@@ -45,7 +46,7 @@ class Config:
 
 class Nexus(list):
     """
-    A NEXUS object implemented as list of tokens with methods to access newick constituents.
+    A NEXUS object implemented as list of commands with methods to read and write blocks.
 
     From the spec:
 
@@ -95,6 +96,7 @@ def __init__(self,
         """
         self.cfg = config or Config(**kw)
         self.trailing_whitespace = []
+        self.leading = []
         self.block_implementations = {}
         for cls in Block.__subclasses__():
             self.block_implementations[cls.__name__.upper()] = cls
@@ -120,7 +122,10 @@ def __init__(self,
                     if token.is_semicolon:
                         commands.append(Command(tuple(tokens)))
                         tokens = []
-            self.trailing_whitespace = tokens
+            if commands:
+                self.trailing_whitespace = tokens
+            else:
+                self.leading = tokens
             s = commands
         list.__init__(self, s)
 
@@ -211,6 +216,7 @@ def __str__(self):
             END;
         """
         return NEXUS \
+            + ''.join(str(t) for t in self.leading) \
             + ''.join(''.join(str(t) for t in cmd) for cmd in self) \
             + ''.join(str(t) for t in self.trailing_whitespace)
 
@@ -224,8 +230,10 @@ def to_file(self, p: typing.Union[str, pathlib.Path]):
         p.write_text(text, encoding=self.cfg.encoding)
 
     def iter_comments(self):
+        yield from (t for t in self.leading if t.type == TokenType.COMMENT)
         for cmd in self:
             yield from (t for t in cmd if t.type == TokenType.COMMENT)
+        yield from (t for t in self.trailing_whitespace if t.type == TokenType.COMMENT)
 
     @property
     def comments(self) -> typing.List[str]:
@@ -262,13 +270,18 @@ def iter_blocks(self):
 
     def validate(self, log=None):
         valid = True
+        if any(t.type not in {TokenType.WHITESPACE, TokenType.COMMENT} for t in self.leading):
+            log_or_raise('Invalid token in preamble', log=log)
         for block in self.iter_blocks():
             #
             # FIXME: we can do a lot of validation here! If block.__commands__ is a list, there is
             # some fixed order between commands.
             # If Payload.__multivalued__ == False, only one command instance is allowed, ...
             #
             valid = valid and block.validate(log=log)
+        if any(t.type not in {TokenType.WHITESPACE, TokenType.COMMENT}
+               for t in self.trailing_whitespace):
+            log_or_raise('Invalid token in text after the last command', log=log)
         return valid
 
     def get_numbers(self, object_name, items):

tests/test_nexus.py

@@ -161,3 +161,18 @@ def test_Booleans_With_Values(fixture_dir):
 def test_Mesquite_multitaxa(fixture_dir):
     nex = Nexus.from_file(fixture_dir / 'multitaxa_mesquite.nex')
     assert [char.get_matrix() for char in nex.blocks['CHARACTERS']]
+
+
+def test_comments():
+    nex = Nexus('#NEXUS\n[comment]')
+    assert nex.comments == ['comment']
+    nex.append_block(Block.from_commands([], name='b', comment='block comment'))
+    assert nex.comments == ['comment', 'block comment']
+
+    with pytest.raises(ValueError):
+        nex = Nexus('#NEXUS\n[comment]]')
+        nex.validate()
+
+    with pytest.raises(ValueError):
+        nex = Nexus('#NEXUS\nBEGIN b; end;[comment]]')
+        nex.validate()