Added iterator classes to simplify parsing the definition files

pint/util.py

@@ -26,6 +26,7 @@ from tokenize import untokenize
 from .compat import string_types, tokenizer, lru_cache, NullHandler, maketrans, NUMERIC_TYPES
 from .formatting import format_unit,siunitx_format_unit
 from .pint_eval import build_eval_tree
+from .errors import DefinitionSyntaxError
 
 logger = logging.getLogger(__name__)
 logger.addHandler(NullHandler())
@@ -653,3 +654,72 @@ def fix_str_conversions(cls):
     else:
         cls.__str__ = __unicode__
     return cls
+
+
+class SourceIterator(object):
+    """Iterator to facilitate reading the definition files.
+
+    Accepts any sequence (like a list of lines, a file or another SourceIterator)
+
+    The iterator yields the line number and line (skipping comments and empty lines)
+    and stripping white spaces.
+
+    for lineno, line in SourceIterator(sequence):
+        # do something here
+
+    """
+
+    def __new__(cls, sequence):
+        if isinstance(sequence, SourceIterator):
+            return sequence
+
+        obj = object.__new__(cls)
+
+        if sequence is not None:
+            obj.internal = enumerate(sequence, 1)
+            obj.last = (None, None)
+
+        return obj
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        line = ''
+        while not line or line.startswith('#'):
+            lineno, line = next(self.internal)
+            line = line.strip()
+
+        self.last = lineno, line
+        return lineno, line
+
+    def block_iter(self):
+        """Iterate block including header.
+        """
+        return BlockIterator(self)
+
+
+class BlockIterator(SourceIterator):
+    """Like SourceIterator but stops when it finds '@end'
+    It also raises an error if another '@' directive is found inside.
+    """
+
+    def __new__(cls, line_iterator):
+        obj = SourceIterator.__new__(cls, None)
+        obj.internal = line_iterator.internal
+        obj.last = line_iterator.last
+        obj.done_last = False
+        return obj
+
+    def __next__(self):
+        if not self.done_last:
+            self.done_last = True
+            return self.last
+
+        lineno, line = SourceIterator.__next__(self)
+        if line.startswith('@end'):
+            raise StopIteration
+        elif line.startswith('@'):
+            raise DefinitionSyntaxError('cannot nest @ directives', lineno=lineno)
+
+        return lineno, line