add options for include directive (#684)

Co-authored-by: Maximilian Hils <git@maximilianhils.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ## Unreleased: pdoc next
 
+- The `.. include:` rST directive now supports start-line, end-line, start-after, end-before options.
+  ([#684](https://github.com/mitmproxy/pdoc/pull/684), @frankharkins)
 - Fix image embedding in included rST files.
   ([#692](https://github.com/mitmproxy/pdoc/pull/692), @meghprkh)
 - Support type-hints from stub-only packages. E.g: `scipy-stubs`

pdoc/__init__.py

@@ -283,7 +283,16 @@ class GoldenRetriever(Dog):
 """
 ```
 
-Since version 11, pdoc processes such reStructuredText elements by default.
+You can also include only parts of a file with the
+[`start-line`, `end-line`, `start-after`, and `end-after` options](https://docutils.sourceforge.io/docs/ref/rst/directives.html#including-an-external-document-fragment):
+
+```python
+"""
+.. include:: ../README.md
+   :start-line: 1
+   :end-before: Changelog
+"""
+```
 
 
 ## ...add a title page?

pdoc/docstrings.py

@@ -360,6 +360,38 @@ def replace_link(m: re.Match[str]) -> str:
     return contents
 
 
+def _rst_extract_options(contents: str) -> tuple[str, dict[str, str]]:
+    """
+    Extract options from the beginning of reStructuredText directives.
+
+    Return the trimmed content and a dict of options.
+    """
+    options = {}
+    while match := re.match(r"^\s*:(.+?):(.*)([\s\S]*)", contents):
+        key, value, contents = match.groups()
+        options[key] = value.strip()
+
+    return contents, options
+
+
+def _rst_include_trim(contents: str, options: dict[str, str]) -> str:
+    """
+    <https://docutils.sourceforge.io/docs/ref/rst/directives.html#include-options>
+    """
+    if "end-line" in options or "start-line" in options:
+        lines = contents.splitlines()
+        if i := options.get("end-line"):
+            lines = lines[: int(i)]
+        if i := options.get("start-line"):
+            lines = lines[int(i) :]
+        contents = "\n".join(lines)
+    if x := options.get("end-before"):
+        contents = contents[: contents.index(x)]
+    if x := options.get("start-after"):
+        contents = contents[contents.index(x) + len(x) :]
+    return contents
+
+
 def _rst_admonitions(contents: str, source_file: Path | None) -> str:
     """
     Convert reStructuredText admonitions - a bit tricky because they may already be indented themselves.
@@ -371,6 +403,7 @@ def _rst_admonition(m: re.Match[str]) -> str:
         type = m.group("type")
         val = m.group("val").strip()
         contents = dedent(m.group("contents")).strip()
+        contents, options = _rst_extract_options(contents)
 
         if type == "include":
             loc = source_file or Path(".")
@@ -379,6 +412,10 @@ def _rst_admonition(m: re.Match[str]) -> str:
             except OSError as e:
                 warnings.warn(f"Cannot include {val!r}: {e}")
                 included = "\n"
+            try:
+                included = _rst_include_trim(included, options) + "\n"
+            except ValueError as e:
+                warnings.warn(f"Failed to process include options for {val!r}: {e}")
             included = _rst_admonitions(included, loc.parent / val)
             included = embed_images(included, loc.parent / val)
             return indent(included, ind)

test/test_docstrings.py

@@ -1,11 +1,16 @@
+from pathlib import Path
+
 from hypothesis import given
 from hypothesis.strategies import text
 import pytest
 
 from pdoc import docstrings
 
 # The important tests are in test_snapshot.py (and, by extension, testdata/)
-# only some fuzzing here.
+# mostly some fuzzing here.
+
+
+here = Path(__file__).parent.absolute()
 
 
 @given(text())
@@ -26,6 +31,59 @@ def test_rst(s):
     assert not s or ret
 
 
+@given(text())
+def test_rst_extract_options_fuzz(s):
+    content, options = docstrings._rst_extract_options(s)
+    assert not s or content or options
+
+
+def test_rst_extract_options():
+    content = (
+        ":alpha: beta\n"
+        ":charlie:delta:foxtrot\n"
+        "rest of content\n"
+        ":option ignored: as follows content\n"
+    )
+    content, options = docstrings._rst_extract_options(content)
+    assert options == {
+        "alpha": "beta",
+        "charlie": "delta:foxtrot",
+    }
+    assert content == ("\nrest of content\n" ":option ignored: as follows content\n")
+
+
+def test_rst_include_trim_lines():
+    content = "alpha\nbeta\ncharlie\ndelta\necho"
+    trimmed = docstrings._rst_include_trim(
+        content, {"start-line": "2", "end-line": "4"}
+    )
+    assert trimmed == "charlie\ndelta"
+
+
+def test_rst_include_trim_pattern():
+    content = "alpha\nbeta\ncharlie\ndelta\necho"
+    trimmed = docstrings._rst_include_trim(
+        content, {"start-after": "beta", "end-before": "echo"}
+    )
+    assert trimmed == "\ncharlie\ndelta\n"
+
+
+def test_rst_include_trim_mixture():
+    content = "alpha\nbeta\ncharlie\ndelta\necho"
+    trimmed = docstrings._rst_include_trim(
+        content, {"start-after": "beta", "end-line": "4"}
+    )
+    assert trimmed == "\ncharlie\ndelta"
+
+
 def test_rst_include_nonexistent():
     with pytest.warns(UserWarning, match="Cannot include 'nonexistent.txt'"):
         docstrings.rst(".. include:: nonexistent.txt", None)
+
+
+def test_rst_include_invalid_options():
+    with pytest.warns(UserWarning, match="Failed to process include options"):
+        docstrings.rst(
+            ".. include:: ../README.md\n   :start-line: invalid",
+            here / "test_docstrings.py",
+        )