Add prepend_prolog() and append_epilog()

3 files changed, 108 insertions, 9 deletions

diff --git a/sphinx/io.py b/sphinx/io.py
index b35dcc73f..88b7eca36 100644
--- a/sphinx/io.py
+++ b/sphinx/io.py
@@ -9,7 +9,6 @@
     :license: BSD, see LICENSE for details.
 """
 import codecs
-import re
 import warnings
 
 from docutils.core import Publisher
@@ -37,6 +36,7 @@ from sphinx.transforms.i18n import (
 from sphinx.transforms.references import SphinxDomains, SubstitutionDefinitionsRemover
 from sphinx.util import logging
 from sphinx.util.docutils import LoggingReporter
+from sphinx.util.rst import append_epilog, docinfo_re, prepend_prolog
 from sphinx.versioning import UIDTransform
 
 if False:
@@ -51,8 +51,6 @@ if False:
     from sphinx.environment import BuildEnvironment  # NOQA
     from sphinx.util.typing import unicode  # NOQA
 
-docinfo_re = re.compile(':\\w+:.*?')
-
 
 logger = logging.getLogger(__name__)
 
@@ -260,10 +258,8 @@ class SphinxRSTFileInput(SphinxBaseFileInput):
         for lineno, line in enumerate(lines):
             content.append(line, self.source_path, lineno)
 
-        if self.env.config.rst_prolog:
-            self.prepend_prolog(content, self.env.config.rst_prolog)
-        if self.env.config.rst_epilog:
-            self.append_epilog(content, self.env.config.rst_epilog)
+        prepend_prolog(content, self.env.config.rst_prolog)
+        append_epilog(content, self.env.config.rst_epilog)
 
         return content
 
diff --git a/sphinx/util/rst.py b/sphinx/util/rst.py
index 8b8aca9ea..1ad7bc468 100644
--- a/sphinx/util/rst.py
+++ b/sphinx/util/rst.py
@@ -24,11 +24,14 @@ from sphinx.util import logging
 if False:
     # For type annotation
     from typing import Generator  # NOQA
+    from docutils.statemachine import StringList  # NOQA
     from sphinx.util.typing import unicode  # NOQA
 
-symbols_re = re.compile(r'([!-\-/:-@\[-`{-~])')  # symbols without dot(0x2e)
 logger = logging.getLogger(__name__)
 
+docinfo_re = re.compile(':\\w+:.*?')
+symbols_re = re.compile(r'([!-\-/:-@\[-`{-~])')  # symbols without dot(0x2e)
+
 
 def escape(text):
     # type: (unicode) -> unicode
@@ -51,3 +54,35 @@ def default_role(docname, name):
     yield
 
     docutils.unregister_role('')
+
+
+def prepend_prolog(content, prolog):
+    # type: (StringList, unicode) -> None
+    """Prepend a string to content body as prolog."""
+    if prolog:
+        pos = 0
+        for line in content:
+            if docinfo_re.match(line):
+                pos += 1
+            else:
+                break
+
+        if pos > 0:
+            # insert a blank line after docinfo
+            content.insert(pos, '', '<generated>', 0)
+            pos += 1
+
+        # insert prolog (after docinfo if exists)
+        for lineno, line in enumerate(prolog.splitlines()):
+            content.insert(pos + lineno, line, '<rst_prolog>', lineno)
+
+        content.insert(pos + lineno + 1, '', '<generated>', 0)
+
+
+def append_epilog(content, epilog):
+    # type: (StringList, unicode) -> None
+    """Append a string to content body as epilog."""
+    if epilog:
+        content.append('', '<generated>', 0)
+        for lineno, line in enumerate(epilog.splitlines()):
+            content.append(line, '<rst_epilog>', lineno)
diff --git a/tests/test_util_rst.py b/tests/test_util_rst.py
index 07b9174cc..166c865d3 100644
--- a/tests/test_util_rst.py
+++ b/tests/test_util_rst.py
@@ -8,7 +8,10 @@
     :copyright: Copyright 2007-2018 by the Sphinx team, see AUTHORS.
     :license: BSD, see LICENSE for details.
 """
-from sphinx.util.rst import escape
+
+from docutils.statemachine import StringList
+
+from sphinx.util.rst import append_epilog, escape, prepend_prolog
 
 
 def test_escape():
@@ -16,3 +19,68 @@ def test_escape():
     assert escape('footnote [#]_') == r'footnote \[\#\]\_'
     assert escape('sphinx.application') == r'sphinx.application'
     assert escape('.. toctree::') == r'\.. toctree\:\:'
+
+
+def test_append_epilog(app):
+    epilog = 'this is rst_epilog\ngood-bye reST!'
+    content = StringList(['hello Sphinx world',
+                          'Sphinx is a document generator'],
+                         'dummy.rst')
+    append_epilog(content, epilog)
+
+    assert list(content.xitems()) == [('dummy.rst', 0, 'hello Sphinx world'),
+                                      ('dummy.rst', 1, 'Sphinx is a document generator'),
+                                      ('<generated>', 0, ''),
+                                      ('<rst_epilog>', 0, 'this is rst_epilog'),
+                                      ('<rst_epilog>', 1, 'good-bye reST!')]
+
+
+def test_prepend_prolog(app):
+    prolog = 'this is rst_prolog\nhello reST!'
+    content = StringList([':title: test of SphinxFileInput',
+                          ':author: Sphinx team',
+                          '',
+                          'hello Sphinx world',
+                          'Sphinx is a document generator'],
+                         'dummy.rst')
+    prepend_prolog(content, prolog)
+
+    assert list(content.xitems()) == [('dummy.rst', 0, ':title: test of SphinxFileInput'),
+                                      ('dummy.rst', 1, ':author: Sphinx team'),
+                                      ('<generated>', 0, ''),
+                                      ('<rst_prolog>', 0, 'this is rst_prolog'),
+                                      ('<rst_prolog>', 1, 'hello reST!'),
+                                      ('<generated>', 0, ''),
+                                      ('dummy.rst', 2, ''),
+                                      ('dummy.rst', 3, 'hello Sphinx world'),
+                                      ('dummy.rst', 4, 'Sphinx is a document generator')]
+
+
+def test_prepend_prolog_with_CR(app):
+    # prolog having CR at tail
+    prolog = 'this is rst_prolog\nhello reST!\n'
+    content = StringList(['hello Sphinx world',
+                          'Sphinx is a document generator'],
+                         'dummy.rst')
+    prepend_prolog(content, prolog)
+
+    assert list(content.xitems()) == [('<rst_prolog>', 0, 'this is rst_prolog'),
+                                      ('<rst_prolog>', 1, 'hello reST!'),
+                                      ('<generated>', 0, ''),
+                                      ('dummy.rst', 0, 'hello Sphinx world'),
+                                      ('dummy.rst', 1, 'Sphinx is a document generator')]
+
+
+def test_prepend_prolog_without_CR(app):
+    # prolog not having CR at tail
+    prolog = 'this is rst_prolog\nhello reST!'
+    content = StringList(['hello Sphinx world',
+                          'Sphinx is a document generator'],
+                         'dummy.rst')
+    prepend_prolog(content, prolog)
+
+    assert list(content.xitems()) == [('<rst_prolog>', 0, 'this is rst_prolog'),
+                                      ('<rst_prolog>', 1, 'hello reST!'),
+                                      ('<generated>', 0, ''),
+                                      ('dummy.rst', 0, 'hello Sphinx world'),
+                                      ('dummy.rst', 1, 'Sphinx is a document generator')]