coalescestreams.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. """Preprocessor for merging consecutive stream outputs for easier handling."""
  2. # Copyright (c) IPython Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. import re
  5. import functools
  6. from traitlets.log import get_logger
  7. def cell_preprocessor(function):
  8. """
  9. Wrap a function to be executed on all cells of a notebook
  10. The wrapped function should have these parameters:
  11. cell : NotebookNode cell
  12. Notebook cell being processed
  13. resources : dictionary
  14. Additional resources used in the conversion process. Allows
  15. preprocessors to pass variables into the Jinja engine.
  16. index : int
  17. Index of the cell being processed
  18. """
  19. @functools.wraps(function)
  20. def wrappedfunc(nb, resources):
  21. get_logger().debug(
  22. "Applying preprocessor: %s", function.__name__
  23. )
  24. for index, cell in enumerate(nb.cells):
  25. nb.cells[index], resources = function(cell, resources, index)
  26. return nb, resources
  27. return wrappedfunc
  28. cr_pat = re.compile(r'.*\r(?=[^\n])')
  29. @cell_preprocessor
  30. def coalesce_streams(cell, resources, index):
  31. """
  32. Merge consecutive sequences of stream output into single stream
  33. to prevent extra newlines inserted at flush calls
  34. Parameters
  35. ----------
  36. cell : NotebookNode cell
  37. Notebook cell being processed
  38. resources : dictionary
  39. Additional resources used in the conversion process. Allows
  40. transformers to pass variables into the Jinja engine.
  41. index : int
  42. Index of the cell being processed
  43. """
  44. outputs = cell.get('outputs', [])
  45. if not outputs:
  46. return cell, resources
  47. last = outputs[0]
  48. new_outputs = [last]
  49. for output in outputs[1:]:
  50. if (output.output_type == 'stream' and
  51. last.output_type == 'stream' and
  52. last.name == output.name
  53. ):
  54. last.text += output.text
  55. else:
  56. new_outputs.append(output)
  57. last = output
  58. # process \r characters
  59. for output in new_outputs:
  60. if output.output_type == 'stream' and '\r' in output.text:
  61. output.text = cr_pat.sub('', output.text)
  62. cell.outputs = new_outputs
  63. return cell, resources