feat: Add trim_doctest_flag to google and numpy parsers

Author: thatlittleboy

src/pytkdocs/parsers/docstrings/google.py

@@ -28,20 +28,26 @@
 
 RE_GOOGLE_STYLE_ADMONITION: Pattern = re.compile(r"^(?P<indent>\s*)(?P<type>[\w-]+):((?:\s+)(?P<title>.+))?$")
 """Regular expressions to match lines starting admonitions, of the form `TYPE: [TITLE]`."""
+RE_DOCTEST_BLANKLINE: Pattern = re.compile(r"^\s*<BLANKLINE>\s*$")
+"""Regular expression to match lines of the form `<BLANKLINE>`."""
+RE_DOCTEST_FLAGS: Pattern = re.compile(r"(\s*#\s*doctest:.+)$")
+"""Regular expression to match lines containing doctest flags of the form `# doctest: +FLAG`."""
 
 
 class Google(Parser):
     """A Google-style docstrings parser."""
 
-    def __init__(self, replace_admonitions: bool = True) -> None:
+    def __init__(self, replace_admonitions: bool = True, trim_doctest_flags: bool = True) -> None:
         """
         Initialize the object.
 
         Arguments:
             replace_admonitions: Whether to replace admonitions by their Markdown equivalent.
+            trim_doctest_flags: Whether to remove doctest flags.
         """
         super().__init__()
         self.replace_admonitions = replace_admonitions
+        self.trim_doctest_flags = trim_doctest_flags
         self.section_reader = {
             Section.Type.PARAMETERS: self.read_parameters_section,
             Section.Type.KEYWORD_ARGS: self.read_keyword_arguments_section,
@@ -492,6 +498,9 @@ def read_examples_section(self, lines: List[str], start_index: int) -> Tuple[Opt
                     current_text.append(line)
 
             elif in_code_example:
+                if self.trim_doctest_flags:
+                    line = RE_DOCTEST_FLAGS.sub("", line)
+                    line = RE_DOCTEST_BLANKLINE.sub("", line)
                 current_example.append(line)
 
             elif line.startswith("```"):
@@ -506,6 +515,9 @@ def read_examples_section(self, lines: List[str], start_index: int) -> Tuple[Opt
                     sub_sections.append((Section.Type.MARKDOWN, "\n".join(current_text)))
                     current_text = []
                 in_code_example = True
+
+                if self.trim_doctest_flags:
+                    line = RE_DOCTEST_FLAGS.sub("", line)
                 current_example.append(line)
 
             else:

src/pytkdocs/parsers/docstrings/numpy.py

@@ -1,21 +1,30 @@
 """This module defines functions and classes to parse docstrings into structured data."""
 import re
-from typing import List, Optional
+from typing import List, Optional, Pattern
 
 from docstring_parser import parse
 from docstring_parser.common import Docstring, DocstringMeta
 
 from pytkdocs.parsers.docstrings.base import AnnotatedObject, Attribute, Parameter, Parser, Section, empty
 
+RE_DOCTEST_BLANKLINE: Pattern = re.compile(r"^\s*<BLANKLINE>\s*$")
+"""Regular expression to match lines of the form `<BLANKLINE>`."""
+RE_DOCTEST_FLAGS: Pattern = re.compile(r"(\s*#\s*doctest:.+)$")
+"""Regular expression to match lines containing doctest flags of the form `# doctest: +FLAG`."""
+
 
 class Numpy(Parser):
     """A Numpy-style docstrings parser."""
 
-    def __init__(self) -> None:
+    def __init__(self, trim_doctest_flags: bool = True) -> None:
         """
         Initialize the objects.
+
+        Arguments:
+            trim_doctest_flags: Whether to remove doctest flags.
         """
         super().__init__()
+        self.trim_doctest_flags = trim_doctest_flags
         self.section_reader = {
             Section.Type.PARAMETERS: self.read_parameters_section,
             Section.Type.EXCEPTIONS: self.read_exceptions_section,
@@ -229,6 +238,9 @@ def read_examples_section(
                         current_text.append(line)
 
                 elif in_code_example:
+                    if self.trim_doctest_flags:
+                        line = RE_DOCTEST_FLAGS.sub("", line)
+                        line = RE_DOCTEST_BLANKLINE.sub("", line)
                     current_example.append(line)
 
                 elif line.startswith("```"):
@@ -243,15 +255,21 @@ def read_examples_section(
                         sub_sections.append((Section.Type.MARKDOWN, "\n".join(current_text)))
                         current_text = []
                     in_code_example = True
+
+                    if self.trim_doctest_flags:
+                        line = RE_DOCTEST_FLAGS.sub("", line)
                     current_example.append(line)
                 else:
                     current_text.append(line)
+
         if current_text:
             sub_sections.append((Section.Type.MARKDOWN, "\n".join(current_text)))
         elif current_example:
             sub_sections.append((Section.Type.EXAMPLES, "\n".join(current_example)))
+
         if sub_sections:
             return Section(Section.Type.EXAMPLES, sub_sections)
+
         if re.search("Examples\n", docstring):
             self.error("Empty examples section")
         return None

tests/test_parsers/test_docstrings/test_google.py

@@ -14,10 +14,19 @@ class DummyObject:
     path = "o"
 
 
-def parse(docstring, signature=None, return_type=inspect.Signature.empty, admonitions=True):
+def parse(
+    docstring,
+    signature=None,
+    return_type=inspect.Signature.empty,
+    admonitions=True,
+    trim_doctest=False,
+):
     """Helper to parse a doctring."""
-    return Google(replace_admonitions=admonitions).parse(
-        dedent(docstring).strip(), {"obj": DummyObject(), "signature": signature, "type": return_type}
+    parser = Google(replace_admonitions=admonitions, trim_doctest_flags=trim_doctest)
+
+    return parser.parse(
+        dedent(docstring).strip(),
+        context={"obj": DummyObject(), "signature": signature, "type": return_type},
     )
 
 
@@ -129,8 +138,48 @@ def f(x: int, y: int, *, z: int) -> int:
     assert not errors
 
 
+def test_function_with_examples_trim_doctest():
+    """Parse example docstring with trim_doctest_flags option."""
+
+    def f(x: int) -> int:
+        """Test function.
+
+        Example:
+
+            We want to skip the following test.
+            >>> 1 + 1 == 3  # doctest: +SKIP
+            True
+
+            And then a few more examples here:
+            >>> print("a\\n\\nb")
+            a
+            <BLANKLINE>
+            b
+            >>> 1 + 1 == 2  # doctest: +SKIP
+            >>> print(list(range(1, 100)))    # doctest: +ELLIPSIS
+            [1, 2, ..., 98, 99]
+        """
+        return x
+
+    sections, errors = parse(
+        inspect.getdoc(f),
+        inspect.signature(f),
+        trim_doctest=True,
+    )
+    assert len(sections) == 2
+    assert len(sections[1].value) == 4
+    assert not errors
+
+    # Verify that doctest flags have indeed been trimmed
+    example_str = sections[1].value[1][1]
+    assert "# doctest: +SKIP" not in example_str
+    example_str = sections[1].value[3][1]
+    assert "<BLANKLINE>" not in example_str
+    assert "\n>>> print(list(range(1, 100)))\n" in example_str
+
+
 def test_function_with_examples():
-    """Parse a function docstring with signature annotations."""
+    """Parse a function docstring with examples."""
 
     def f(x: int, y: int) -> int:
         """

tests/test_parsers/test_docstrings/test_numpy.py

@@ -11,11 +11,18 @@ class DummyObject:
     path = "o"
 
 
-def parse(docstring, signature=None, return_type=inspect.Signature.empty):
+def parse(
+    docstring,
+    signature=None,
+    return_type=inspect.Signature.empty,
+    trim_doctest=False,
+):
     """Helper to parse a doctring."""
-    return Numpy().parse(
+    parser = Numpy(trim_doctest_flags=trim_doctest)
+
+    return parser.parse(
         dedent(docstring).strip(),
-        {"obj": DummyObject(), "signature": signature, "type": return_type},
+        context={"obj": DummyObject(), "signature": signature, "type": return_type},
     )
 
 
@@ -48,13 +55,13 @@ def test_sections_without_signature():
 
         Parameters
         ----------
-        void : 
+        void :
             SEGFAULT.
-        niet : 
+        niet :
             SEGFAULT.
-        nada : 
+        nada :
             SEGFAULT.
-        rien : 
+        rien :
             SEGFAULT.
 
         Raises
@@ -63,7 +70,7 @@ def test_sections_without_signature():
             when nothing works as expected.
 
         Returns
-        ------- 
+        -------
         bool
             Itself.
         """
@@ -135,8 +142,48 @@ def f(x: int, y: int) -> int:
     assert not errors
 
 
+def test_function_with_examples_trim_doctest():
+    """Parse example docstring with trim_doctest_flags option."""
+
+    def f(x: int) -> int:
+        """Test function.
+
+        Example
+        -------
+        We want to skip the following test.
+        >>> 1 + 1 == 3  # doctest: +SKIP
+        True
+
+        And then a few more examples here:
+        >>> print("a\\n\\nb")
+        a
+        <BLANKLINE>
+        b
+        >>> 1 + 1 == 2  # doctest: +SKIP
+        >>> print(list(range(1, 100)))    # doctest: +ELLIPSIS
+        [1, 2, ..., 98, 99]
+        """
+        return x
+
+    sections, errors = parse(
+        inspect.getdoc(f),
+        inspect.signature(f),
+        trim_doctest=True,
+    )
+    assert len(sections) == 2
+    assert len(sections[1].value) == 4
+    assert not errors
+
+    # Verify that doctest flags have indeed been trimmed
+    example_str = sections[1].value[1][1]
+    assert "# doctest: +SKIP" not in example_str
+    example_str = sections[1].value[3][1]
+    assert "<BLANKLINE>" not in example_str
+    assert "\n>>> print(list(range(1, 100)))\n" in example_str
+
+
 def test_function_with_examples():
-    """Parse a function docstring with signature annotations."""
+    """Parse a function docstring with examples."""
 
     def f(x: int, y: int) -> int:
         """