stripped down handler interface

Author: xrotwang

src/nexus/handlers/__init__.py

@@ -19,13 +19,14 @@ class GenericHandler(object):
 
     Handlers have (at least) the following attributes:
 
-        1. parse(self, data) - the function for parsing the block
-        2. write(self, data) - a function for returning the block to a text
+        1. __init__(self, **kw) - the function for parsing the block
+        2. iter_lines(self) - a function for returning the block to a text
             representation (used to regenerate a nexus file).
         3. block - a list of raw strings in this block
     """
-    def __init__(self, data=None):
+    def __init__(self, name=None, data=None):
         """Initialise datastore in <block> under <keyname>"""
+        self.name = name
         self.block = data or []
         self.comments = []
 
@@ -34,13 +35,22 @@ def __init__(self, data=None):
             if line.strip().startswith("[") and line.strip().endswith("]"):
                 self.comments.append(line)
 
+    def iter_lines(self):
+        for i, line in enumerate(self.block):
+            if (i == 0 and BEGIN_PATTERN.search(line)) or \
+                    (i == len(self.block) - 1 and END_PATTERN.search(line)):
+                continue
+            yield line
+
     def write(self):
         """
-        Generates a string containing a generic nexus block for this data.
-
-        :return: String
+        Generates a string containing a nexus block.
         """
-        return "\n".join(self.block)
+        return "".join(
+            ['begin {0};\n'.format(self.name)] +
+            [l + '\n' for l in self.iter_lines()] +
+            ['end;\n']
+        )
 
     @staticmethod
     def remove_comments(line):
@@ -64,11 +74,5 @@ def remove_comments(line):
     def is_mesquite_attribute(line):
         """
         Returns True if the line is a mesquite attribute
-
-        :return: Boolean
         """
-        if MESQUITE_TITLE_PATTERN.match(line):
-            return True
-        elif MESQUITE_LINK_PATTERN.match(line):
-            return True
-        return False
+        return bool(MESQUITE_TITLE_PATTERN.match(line)) or bool(MESQUITE_LINK_PATTERN.match(line))

src/nexus/handlers/data.py

@@ -21,8 +21,8 @@ class DataHandler(GenericHandler):
         re.IGNORECASE | re.DOTALL | re.MULTILINE
     )
 
-    def __init__(self, data=None):
-        super(DataHandler, self).__init__(data)
+    def __init__(self, **kw):
+        super(DataHandler, self).__init__(**kw)
         self.charlabels = {}
         self.attributes = []
         self.format = {}
@@ -118,7 +118,7 @@ def characters(self):
         return self._characters
 
     def is_missing_or_gap(self, state):
-        return True if state in ('-', '?') else False
+        return state in ('-', '?')
 
     def parse_format_line(self, data):
         """
@@ -136,11 +136,10 @@ def parse_format_line(self, data):
         except IndexError:
             return None
 
-        line = line.lower()
-        line = line.replace(" =", "=").replace("= ", "=")  # standardise
+        line = line.lower().replace(" =", "=").replace("= ", "=")  # standardise
         for chunk in WHITESPACE_PATTERN.split(line):
             try:
-                key, value = chunk.split("=")
+                key, value = chunk.split("=", maxsplit=1)
                 value = QUOTED_PATTERN.sub('\\1', value)
             except ValueError:
                 key, value = chunk, True
@@ -240,7 +239,7 @@ def _parse_charstate_block(self, data):
                 char_index += 1
         return new_data.split("\n")
 
-    def write(self):
+    def iter_lines(self):
         """
         Generates a string containing a nexus data block.
 
@@ -269,34 +268,23 @@ def _make_format_line(self):
                     fstring.append("%s=%s" % (key, value))
             return " ".join(fstring) + ";"
 
-        out = []
-        out.append('begin data;')
         for att in self.attributes:
-            out.append("\t%s" % att)
-        out.append('\tdimensions ntax=%d nchar=%d;' % (self.ntaxa, self.nchar))
-        out.append(_make_format_line(self))
+            yield "\t%s" % att
+        yield '\tdimensions ntax=%d nchar=%d;' % (self.ntaxa, self.nchar)
+        yield _make_format_line(self)
         # handle char block
         if self.charlabels:
-            out.append('\tcharstatelabels')
+            yield '\tcharstatelabels'
             max_id = max(self.charlabels)
             for char_id in sorted(self.charlabels):
-                out.append('\t\t%d %s%s' % (
-                    char_id + 1,  # zero-indexing
-                    self.charlabels[char_id],
-                    ',' if char_id < max_id else ''
-                ))
-            out.append('\t;')
-        out.append("matrix")
+                yield '\t\t%d %s%s' % (  # zero-indexing
+                    char_id + 1, self.charlabels[char_id], ',' if char_id < max_id else '')
+            yield '\t;'
+        yield "matrix"
         max_taxon_len = max([len(_) for _ in self.matrix])
         for taxon in sorted(self.matrix):
-            out.append("%s %s" % (taxon.ljust(max_taxon_len), ''.join(self.matrix[taxon])))
-        out.append(" ;")
-        out.append("end;")
-        return "\n".join(out)
-
-    def __repr__(self):
-        return "<NexusDataBlock: %d characters from %d taxa>" % \
-            (self.nchar, self.ntaxa)
+            yield "%s %s" % (taxon.ljust(max_taxon_len), ''.join(self.matrix[taxon]))
+        yield " ;"
 
 
 class CharacterHandler(DataHandler):

src/nexus/handlers/taxa.py

@@ -13,8 +13,8 @@ class TaxaHandler(GenericHandler):
     is_dimensions = re.compile(r"""dimensions\s*ntax\s*=\s*(\d+)""", re.IGNORECASE)
     is_taxlabel_block = re.compile(r"""\btaxlabels\b""", re.IGNORECASE)
 
-    def __init__(self, data=None):
-        super(TaxaHandler, self).__init__(data)
+    def __init__(self, **kw):
+        super(TaxaHandler, self).__init__(**kw)
         self.taxa = []
         self.attributes = []
         self.annotations = {}
@@ -69,28 +69,16 @@ def _parse_taxa(self, line):
             if taxon and taxon not in self.taxa:
                 yield (taxon, annot)
 
-    def write(self):
-        """
-        Generates a string containing a taxa block for this data.
-
-        :return: String
-        """
+    def iter_lines(self):
         def wrap(s):
             return s if ' ' not in s else "'%s'" % s
 
-        out = ['begin taxa;']
-        # handle any attributes
         for att in self.attributes:
-            out.append("\t%s" % att)
-        out.append('\tdimensions ntax=%d;' % self.ntaxa)
-        out.append('\ttaxlabels')
+            yield "\t%s" % att
+        yield '\tdimensions ntax=%d;' % self.ntaxa
+        yield '\ttaxlabels'
         # taxa labels
         for idx, taxon in enumerate(self.taxa, 1):
             taxon = "%s%s" % (taxon, self.annotations.get(taxon, ''))
-            out.append("\t[%d] %s" % (idx, wrap(taxon)))
-        out.append(';')
-        out.append('end;')
-        return "\n".join(out)
-
-    def __repr__(self):
-        return "<NexusTaxaBlock: %d taxa>" % self.ntaxa
+            yield "\t[%d] %s" % (idx, wrap(taxon))
+        yield ';'

src/nexus/handlers/tree.py

@@ -48,8 +48,8 @@ class TreeHandler(GenericHandler):
         (?=[),])?           # end boundary
     """, re.IGNORECASE + re.VERBOSE + re.DOTALL)
 
-    def __init__(self, data=None):
-        super(TreeHandler, self).__init__(data)
+    def __init__(self, **kw):
+        super(TreeHandler, self).__init__(**kw)
         # does the treefile have a translate block?
         self.was_translated = False
         # has detranslate been called?
@@ -115,12 +115,13 @@ def detranslate(self):
             self.trees[idx] = Tree(self._detranslate_tree(tree, self.translators))
         self._been_detranslated = True
 
-    def _findall_chunks(self, tree):
+    @staticmethod
+    def _findall_chunks(tree):
         """Helper function to find groups used by detranslate."""
         matches = []
         index = 0
         while True:
-            match = self.translate_regex.search(tree, index)
+            match = TreeHandler.translate_regex.search(tree, index)
             if not match:
                 break
             m = dict(zip(['start', 'taxon', 'comment', 'branch'], match.groups()))
@@ -163,27 +164,16 @@ def _detranslate_tree(self, tree, translatetable):
                 tree = tree.replace(found['match'], sub)
         return tree
 
-    def write(self):
-        """
-        Generates a string containing a trees block.
-
-        :return: String
-        """
-        out = ['begin trees;']
+    def iter_lines(self):
         for attr in self.attributes:
-            out.append("\t" + attr)
+            yield "\t" + attr
         if self.was_translated and not self._been_detranslated:
-            out.append('\ttranslate')
-            for index in sorted([int(k) for k in self.translators.keys()]):
-                out.append("\t%d %s," % (index, self.translators[str(index)]))
-            # handle last taxa label in translate block
-            out[-1] = out[-1].replace(',', '')
+            yield '\ttranslate'
+            translator_keys = [int(k) for k in self.translators.keys()]
+            for i, index in enumerate(sorted(translator_keys), start=1):
+                yield "\t%d %s%s" % (
+                    index, self.translators[str(index)], '' if i == len(translator_keys) else ',')
             # work around bug https://github.com/CompEvol/beast2/issues/713
-            out.append(';')
+            yield ';'
         for tree in self.trees:
-            out.append("\t" + tree)
-        out.append('end;\n')
-        return "\n".join(out)
-
-    def __repr__(self):
-        return "<NexusTreeBlock: %d trees>" % self.ntrees
+            yield "\t" + tree

src/nexus/reader.py

@@ -75,7 +75,8 @@ def _set_blocks(self, blocks):
         for block, lines in (blocks.items() if isinstance(blocks, dict) else blocks):
             if block in self.blocks:
                 raise NexusFormatException("Duplicate Block %s" % block)
-            self.blocks[block] = HANDLERS.get(block, GenericHandler)(lines)
+            self.blocks[block] = HANDLERS.get(block, GenericHandler)(
+                name='data' if block == 'characters' else block, data=lines)
 
         if self.blocks.get('characters') and not self.blocks.get('data'):
             self.blocks['data'] = self.blocks['characters']

tests/test_handler_DataHandler.py

@@ -30,9 +30,6 @@ def test_DataHandler_parse_sites(input, expected):
     'Louise': ['1', '1'],
 }
 
-def test_repr(nex):
-    assert repr(nex.data) == "<NexusDataBlock: 2 characters from 4 taxa>"
-
 def test_block_find(nex):
     assert 'data' in nex.blocks
     assert hasattr(nex, 'data')

tests/test_handler_GenericHandler.py

@@ -15,8 +15,9 @@ def test_generic_readwrite():
     ]
     nex = NexusReader.from_string("\n".join(expected))
     for line in nex.sets.write().split("\n"):
-        e = expected.pop(0).strip()
-        assert line.strip() == e
+        if line:
+            e = expected.pop(0).strip()
+            assert line.strip() == e
 
 
 def test_write_produces_end():

tests/test_handler_TaxaHandler.py

@@ -22,10 +22,6 @@ def test_iterable(nex2):
         assert taxa == expected[idx]
 
 
-def test_repr(nex2):
-    assert repr(nex2.blocks['taxa']) == "<NexusTaxaBlock: 4 taxa>"
-
-
 def test_wrap_label_in_quotes_only_when_needed(nex2):
     nex2.taxa.taxa[0] = "long name"
     output = nex2.taxa.write()

tests/test_handler_TreeHandler.py

@@ -42,10 +42,6 @@ def test_write(trees, examples):
     assert expected == written
 
 
-def test_repr(trees):
-    assert repr(trees.trees) == '<NexusTreeBlock: 3 trees>'
-
-
 def test_write_produces_end(trees):
     assert "end;" in trees.trees.write()
     assert len([_ for _ in trees.trees[0].newick_tree.walk()]) == 25