Implement configurable section header adornments (#120)

anjos · LilSpazJoekp · web-flow · commit bd12d34d47d7 · 2025-10-21T10:59:32.000-05:00
Implement configurable section header adornments.

Co-authored-by: LilSpazJoekp &lt;15524072+LilSpazJoekp@users.noreply.github.com&gt;
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -8,7 +8,17 @@ Unreleased
 
 - Added support for Python 3.14.
 - Added ability to skip formatting Python code blocks with the
-  ``--format-python-code-blocks/--no-format-python-code-blocks`` or ``-N`` flag.
+  ``--format-python-code-blocks/--no-format-python-code-blocks`` or ``-N`` command-line
+  argument.
+- Added command-line option ``--section-adornments`` or ``-s`` to control section header
+  adornments.
+- Added command-line argument ``--preserve-adornments`` or ``-pA`` to preserve
+  user-defined section adornments.
+
+**Changed**
+
+- Changed default section heading adornments to match the ones described at
+  https://devguide.python.org/documentation/markup/#sections.
 
 **Fixed**
 
diff --git a/docstrfmt/const.py b/docstrfmt/const.py
@@ -18,7 +18,11 @@
     "**/build",
     "**/dist",
 ]
-SECTION_CHARS = "=-~+.'\"`^_*:#"
+# Part/Chapter/Section adornment characters. The special `|` character separated
+# sections without overlines. If that is not present, then we consider all sections to
+# only contain underlines. From:
+# https://devguide.python.org/documentation/markup/#sections
+SECTION_CHARS = "#*|=-^\"'~+.`_:"
 ROLE_ALIASES = {
     "pep": "PEP",
     "pep-reference": "PEP",
diff --git a/docstrfmt/docstrfmt.py b/docstrfmt/docstrfmt.py
@@ -19,6 +19,7 @@
 
 import black
 import docutils
+import docutils.nodes
 from black import Mode
 from blib2to3.pgen2.tokenize import TokenError
 from docutils import nodes
@@ -29,7 +30,6 @@
 from docutils.utils import Reporter, new_document, unescape
 
 from . import rst_extras
-from .const import SECTION_CHARS
 from .exceptions import InvalidRstError, InvalidRstErrors
 from .rst_extras import _add_directive, generic_role
 from .util import make_enumerator
@@ -151,6 +151,7 @@ def __init__(
         self.ordinal_format = "arabic"
         self.section_depth = 0
         self.subsequent_indent = 0
+        self.use_adornments = None
         for key, value in kwargs.items():
             setattr(self, key, value)
 
@@ -243,6 +244,7 @@ def __init__(
         black_config: Mode | None = None,
         docstring_trailing_line: bool = True,
         format_python_code_blocks: bool = True,
+        section_adornments: list[tuple[str, bool]] | None = None,
     ):
         """Initialize the manager."""
         rst_extras.register()
@@ -262,6 +264,7 @@ def __init__(
         self.docstring_trailing_line = docstring_trailing_line
         self.format_python_code_blocks = format_python_code_blocks
         self._in_docstring = False  # for resolving line numbers in code blocks
+        self.section_adornments = section_adornments
 
     def _patch_unknown_directives(self, text: str) -> None:
         """Patch unknown directives and roles into the parser."""
@@ -365,6 +368,46 @@ def format_node(
         )
         return f"{formatted_node}\n"
 
+    def _register_adornments(
+        self, input_lines: list[str], document: nodes.document
+    ) -> None:
+        """Register adornments from source text on all individual sections.
+
+        This method will parse the document tree and original text to-be-formatted, and
+        will register, at the document tree, the current document configuration
+        representing the adornments for parts, chapters and sections on each level of
+        the document. In particular, it will install an attribute called
+        ``adornment-character`` with the character used for underline or overlining the
+        section, and ``adornment-overline``, if the section should be overlined or not.
+
+        input_lines
+            The lines of input (splitted by newline), that we must format.
+
+        document
+            The pre-parsed document tree, that will be modified with new section
+            attributes as described above.
+
+        """
+        for section in document.traverse(nodes.section):
+            title_node = section.next_node(nodes.title)
+            if (
+                title_node
+                and hasattr(title_node, "line")
+                and title_node.line is not None
+            ):
+                underline = input_lines[title_node.line - 1].strip()[0]
+                overline_lineno = title_node.line - 3
+                overline = False
+
+                if overline_lineno >= 0:
+                    candidate_overline = input_lines[overline_lineno].strip()
+                    if candidate_overline and candidate_overline[0] == underline:
+                        overline = True
+
+                # Store this information in the document tree
+                section["adornment-character"] = underline
+                section["adornment-overline"] = overline
+
     def get_code_line(self, code: str, strict: bool = False) -> int:
         """Get the line number of the code in the file."""
         lines = self.original_text.splitlines()
@@ -407,7 +450,9 @@ def parse_string(
             "", self.settings.report_level, self.settings.halt_level
         )
         parser.parse(text, doc)
-        self._pre_process(doc, line_offset, len(text.splitlines()))
+        input_lines = text.splitlines()
+        self._pre_process(doc, line_offset, len(input_lines))
+        self._register_adornments(input_lines, doc)
         return doc
 
     def perform_format(
@@ -1643,9 +1688,31 @@ def title(
                 None, chain(self._format_children(node, context)), context, node.line
             )
         )
-        char = SECTION_CHARS[context.section_depth - 1]
-        yield text
-        yield char * len(text)
+        char: str = node.parent["adornment-character"]
+        overline: bool = node.parent["adornment-overline"]
+        if context.manager.section_adornments is not None:
+            try:
+                char, overline = context.manager.section_adornments[
+                    context.section_depth - 1
+                ]
+            except IndexError:
+                context.manager.reporter.error(
+                    f"Section at line {node.line} is at depth "
+                    f"{context.section_depth}, however there are only "
+                    f"{len(context.manager.section_adornments)} adornments to pick "
+                    "from. You must review your inputs or change settings."
+                )
+                raise
+
+        if overline:
+            # section headings with overline are centered
+            yield char * (2 + len(text))
+            yield " " + text
+            yield char * (2 + len(text))
+        else:
+            # sections headings without overline are justified
+            yield text
+            yield char * len(text)
 
     def title_reference(
         self,
diff --git a/docstrfmt/main.py b/docstrfmt/main.py
@@ -4,6 +4,7 @@
 
 import asyncio
 import glob
+import itertools
 import logging
 import os
 import signal
@@ -38,6 +39,7 @@
 from libcst.metadata import ParentNodeProvider, PositionProvider
 
 from . import DEFAULT_EXCLUDE, __version__
+from .const import SECTION_CHARS
 from .debug import dump_node
 from .docstrfmt import Manager
 from .exceptions import InvalidRstErrors
@@ -62,12 +64,17 @@ def _format_file(
     mode: Mode,
     docstring_trailing_line: bool,
     format_python_code_blocks: bool,
+    section_adornments: list[tuple[str, bool]] | None,
     raw_output: bool,
     lock: Lock,
 ):
     error_count = 0
     manager = Manager(
-        reporter, mode, docstring_trailing_line, format_python_code_blocks
+        reporter,
+        mode,
+        docstring_trailing_line,
+        format_python_code_blocks,
+        section_adornments,
     )
     if file.name == "-":
         raw_output = True
@@ -290,13 +297,32 @@ def _resolve_length(context: click.Context, _: click.Parameter, value: int | Non
     return value or pyproject_line_length
 
 
+def _validate_adornments(
+    context: click.Context, _: click.Parameter, value: str | None
+) -> list[tuple[str, bool]] | None:
+    actual_value = value or context.params["section_adornments"]
+
+    if len(actual_value) != len(set(actual_value)):
+        msg = "Section adornments must be unique"
+        raise click.BadParameter(msg)
+
+    if "|" in actual_value:
+        with_overline, without_overline = actual_value.split("|", 1)
+        return list(zip(with_overline, itertools.repeat(True))) + list(
+            zip(without_overline, itertools.repeat(False))
+        )
+
+    return list(zip(actual_value, itertools.repeat(False)))
+
+
 async def _run_formatter(
     check: bool,
     file_type: str,
     files: list[str],
     include_txt: bool,
     docstring_trailing_line: bool,
     format_python_code_blocks: bool,
+    section_adornments: list[tuple[str, bool]] | None,
     mode: Mode,
     line_length: int,
     raw_output: bool,
@@ -324,6 +350,7 @@ async def _run_formatter(
                 mode,
                 docstring_trailing_line,
                 format_python_code_blocks,
+                section_adornments,
                 raw_output,
                 lock,
             )
@@ -484,7 +511,25 @@ def leave_FunctionDef(
         self._object_names.pop(-1)
         return updated_node
 
-    def leave_SimpleString(
+    def _escape_quoting(self, node: SimpleString) -> SimpleString:
+        """Escapes quotes in a docstring when necessary."""
+        # handles quoting escaping once
+        for quote in ('"', "'"):
+            quoting = quote * 3
+            if node.value.startswith(quoting) and node.value.endswith(quoting):
+                inner_value = node.value[len(quoting) : -len(quoting)]
+                if quoting in inner_value:
+                    node = node.with_changes(
+                        value=quoting
+                        + inner_value.replace(quoting, f"\\{quoting}").replace(
+                            quoting + quote, f"{quoting}\\{quote}"
+                        )
+                        + quoting
+                    )
+                break
+        return node
+
+    def leave_SimpleString(  # noqa: N802
         self, original_node: SimpleString, updated_node: SimpleString
     ) -> SimpleString:
         """Format the docstring."""
@@ -551,7 +596,7 @@ def leave_SimpleString(
                 self._last_assign = None
                 self._object_names.pop(-1)
                 self._object_type = old_object_type
-        return updated_node
+        return self._escape_quoting(updated_node)
 
     def visit_AssignTarget_target(self, node: AssignTarget) -> None:
         """Set the last assign node."""
@@ -584,7 +629,7 @@ def visit_Module(self, node: Module) -> bool | None:
     is_flag=True,
     help=(
         "Check files and returns a non-zero code if files are not formatted correctly."
-        " Useful for linting. Ignored if raw-input, raw-output, stdin is used."
+        " Useful for linting. Ignored if --raw-input, --raw-output, or stdin is used."
     ),
 )
 @click.option(
@@ -650,6 +695,12 @@ def visit_Module(self, node: Module) -> bool | None:
     ),
     callback=_resolve_length,
 )
+@click.option(
+    "-pA",
+    "--preserve-adornments",
+    help="Preserve existing section adornments.",
+    is_flag=True,
+)
 @click.option(
     "-p",
     "--pyproject-config",
@@ -685,7 +736,26 @@ def visit_Module(self, node: Module) -> bool | None:
     type=str,
 )
 @click.option(
-    "-o", "--raw-output", is_flag=True, help="Output the formatted text to stdout."
+    "-o",
+    "--raw-output",
+    help="Output the formatted text to stdout.",
+    is_flag=True,
+)
+@click.option(
+    "-s",
+    "--section-adornments",
+    type=str,
+    default=SECTION_CHARS,
+    show_default=True,
+    help=(
+        "Define adornments for part/chapter/section headers. It defines a sequence of"
+        " adornments to use for each individual section depth. The list must be"
+        " composed of at least N **distinct** characters for documents with N section"
+        " depths. Provide more if unsure. If the special character `|` (pipe) is"
+        " used, then it defines sections (left portion) that will have overlines"
+        " besides underlines only (right portion). Overrides --preserve-adornments."
+    ),
+    callback=_validate_adornments,
 )
 @click.option(
     "-v",
@@ -710,10 +780,12 @@ def main(
     ignore_cache: bool,
     include_txt: bool,
     line_length: int,
+    preserve_adornments: bool,
     mode: Mode,
     quiet: bool,
     raw_input: str,
     raw_output: bool,
+    section_adornments: list[tuple[str, bool]],
     verbose: int,
     files: list[str],
 ) -> None:
@@ -732,9 +804,17 @@ def main(
         else:
             line_length = DEFAULT_LINE_LENGTH
     error_count = 0
+
+    if preserve_adornments:
+        section_adornments = None
+
     if raw_input:
         manager = Manager(
-            reporter, mode, docstring_trailing_line, format_python_code_blocks
+            reporter,
+            mode,
+            docstring_trailing_line,
+            format_python_code_blocks,
+            section_adornments,
         )
         file = "<raw_input>"
         check = False
@@ -769,6 +849,7 @@ def main(
                 mode,
                 docstring_trailing_line,
                 format_python_code_blocks,
+                section_adornments,
                 raw_output,
                 None,
             )
@@ -800,6 +881,7 @@ def main(
                     include_txt,
                     docstring_trailing_line,
                     format_python_code_blocks,
+                    section_adornments,
                     mode,
                     line_length,
                     raw_output,
diff --git a/tests/test_files/test_file.rst b/tests/test_files/test_file.rst
diff --git a/tests/test_main.py b/tests/test_main.py
