Synnefo » snf-ganeti » ganeti-local

import collections

class LineSplitter:

  """Splits data chunks into lines separated by newline.

  Instances provide a file-like interface.

  """

  def __init__(self, line_fn, *args):

    """Initializes this class.

    @type line_fn: callable

    @param line_fn: Function called for each line, first parameter is line

    @param args: Extra arguments for L{line_fn}

    """

    assert callable(line_fn)

    if args:

      # Python 2.4 doesn't have functools.partial yet

      self._line_fn = \

        lambda line: line_fn(line, *args) # pylint: disable-msg=W0142

    else:

      self._line_fn = line_fn

    self._lines = collections.deque()

    self._buffer = ""

  def write(self, data):

    parts = (self._buffer + data).split("\n")

    self._buffer = parts.pop()

    self._lines.extend(parts)

  def flush(self):

    while self._lines:

      self._line_fn(self._lines.popleft().rstrip("\r\n"))

  def close(self):

    self.flush()

    if self._buffer:

      self._line_fn(self._buffer)

class TestLineSplitter(unittest.TestCase):

  def test(self):

    lines = []

    ls = utils.LineSplitter(lines.append)

    ls.write("Hello World\n")

    self.assertEqual(lines, [])

    ls.write("Foo\n Bar\r\n ")

    ls.write("Baz")

    ls.write("Moo")

    self.assertEqual(lines, [])

    ls.flush()

    self.assertEqual(lines, ["Hello World", "Foo", " Bar"])

    ls.close()

    self.assertEqual(lines, ["Hello World", "Foo", " Bar", " BazMoo"])

  def _testExtra(self, line, all_lines, p1, p2):

    self.assertEqual(p1, 999)

    self.assertEqual(p2, "extra")

    all_lines.append(line)

  def testExtraArgsNoFlush(self):

    lines = []

    ls = utils.LineSplitter(self._testExtra, lines, 999, "extra")

    ls.write("\n\nHello World\n")

    ls.write("Foo\n Bar\r\n ")

    ls.write("")

    ls.write("Baz")

    ls.write("Moo\n\nx\n")

    self.assertEqual(lines, [])

    ls.close()

    self.assertEqual(lines, ["", "", "Hello World", "Foo", " Bar", " BazMoo",

                             "", "x"])