StringIOTree.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. r"""
  2. Implements a buffer with insertion points. When you know you need to
  3. "get back" to a place and write more later, simply call insertion_point()
  4. at that spot and get a new StringIOTree object that is "left behind".
  5. EXAMPLE:
  6. >>> a = StringIOTree()
  7. >>> _= a.write('first\n')
  8. >>> b = a.insertion_point()
  9. >>> _= a.write('third\n')
  10. >>> _= b.write('second\n')
  11. >>> a.getvalue().split()
  12. ['first', 'second', 'third']
  13. >>> c = b.insertion_point()
  14. >>> d = c.insertion_point()
  15. >>> _= d.write('alpha\n')
  16. >>> _= b.write('gamma\n')
  17. >>> _= c.write('beta\n')
  18. >>> b.getvalue().split()
  19. ['second', 'alpha', 'beta', 'gamma']
  20. >>> i = StringIOTree()
  21. >>> d.insert(i)
  22. >>> _= i.write('inserted\n')
  23. >>> out = StringIO()
  24. >>> a.copyto(out)
  25. >>> out.getvalue().split()
  26. ['first', 'second', 'alpha', 'inserted', 'beta', 'gamma', 'third']
  27. """
  28. from __future__ import absolute_import #, unicode_literals
  29. try:
  30. # Prefer cStringIO since io.StringIO() does not support writing 'str' in Py2.
  31. from cStringIO import StringIO
  32. except ImportError:
  33. from io import StringIO
  34. class StringIOTree(object):
  35. """
  36. See module docs.
  37. """
  38. def __init__(self, stream=None):
  39. self.prepended_children = []
  40. if stream is None:
  41. stream = StringIO()
  42. self.stream = stream
  43. self.write = stream.write
  44. self.markers = []
  45. def getvalue(self):
  46. content = [x.getvalue() for x in self.prepended_children]
  47. content.append(self.stream.getvalue())
  48. return "".join(content)
  49. def copyto(self, target):
  50. """Potentially cheaper than getvalue as no string concatenation
  51. needs to happen."""
  52. for child in self.prepended_children:
  53. child.copyto(target)
  54. stream_content = self.stream.getvalue()
  55. if stream_content:
  56. target.write(stream_content)
  57. def commit(self):
  58. # Save what we have written until now so that the buffer
  59. # itself is empty -- this makes it ready for insertion
  60. if self.stream.tell():
  61. self.prepended_children.append(StringIOTree(self.stream))
  62. self.prepended_children[-1].markers = self.markers
  63. self.markers = []
  64. self.stream = StringIO()
  65. self.write = self.stream.write
  66. def insert(self, iotree):
  67. """
  68. Insert a StringIOTree (and all of its contents) at this location.
  69. Further writing to self appears after what is inserted.
  70. """
  71. self.commit()
  72. self.prepended_children.append(iotree)
  73. def insertion_point(self):
  74. """
  75. Returns a new StringIOTree, which is left behind at the current position
  76. (it what is written to the result will appear right before whatever is
  77. next written to self).
  78. Calling getvalue() or copyto() on the result will only return the
  79. contents written to it.
  80. """
  81. # Save what we have written until now
  82. # This is so that getvalue on the result doesn't include it.
  83. self.commit()
  84. # Construct the new forked object to return
  85. other = StringIOTree()
  86. self.prepended_children.append(other)
  87. return other
  88. def allmarkers(self):
  89. children = self.prepended_children
  90. return [m for c in children for m in c.allmarkers()] + self.markers