refactor: move transform logic in ast module; add tests

Author: machow

quartodoc/ast.py

@@ -1,12 +1,109 @@
+from enum import Enum
+from dataclasses import dataclass
 from griffe.docstrings import dataclasses as ds
 from griffe import dataclasses as dc
 from plum import dispatch
 from typing import Union
 
-# TODO: these classes are created to wrap some tuple outputs
-#       we should consolidate logic for transforming the griffe
-#       docstring here (or open a griffe issue).
-from .renderers import tuple_to_data, docstring_section_narrow, ExampleCode, ExampleText
+
+# Transform and patched-in classes ============================================
+# TODO: annotate transform return types. make sure subtypes inherit from correct
+# griffe base objects.
+# TODO: it seems like transform should happen on the root, not individual elements.
+
+
+def transform(el):
+    """Return a more specific docstring element, or simply return the original one."""
+
+    if isinstance(el, tuple):
+        try:
+            return tuple_to_data(el)
+        except ValueError:
+            pass
+    elif isinstance(el, ds.DocstringSection):
+        return _DocstringSectionPatched.transform(el)
+
+    return el
+
+
+# Patch DocstringSection ------------------------------------------------------
+
+
+class DocstringSectionKindPatched(Enum):
+    see_also = "see also"
+    notes = "notes"
+    warnings = "warnings"
+
+
+class _DocstringSectionPatched(ds.DocstringSection):
+    _registry: "dict[Enum, _DocstringSectionPatched]" = {}
+
+    def __init__(self, value: str, title: "str | None" = None):
+        self.value = value
+        super().__init__(title)
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+
+        if cls.kind.value in cls._registry:
+            raise KeyError(f"A section for kind {cls.kind} already exists")
+
+        cls._registry[cls.kind] = cls
+
+    @classmethod
+    def transform(cls, el: ds.DocstringSection) -> ds.DocstringSection:
+        """Attempt to cast DocstringSection element to more specific section type.
+
+        Note that this is meant to patch cases where the general DocstringSectionText
+        class represents a section like See Also, etc..
+        """
+
+        if isinstance(el, ds.DocstringSectionText):
+            for kind, sub_cls in cls._registry.items():
+                prefix = kind.value.title() + "\n---"
+                if el.value.lstrip("\n").startswith(prefix):
+                    stripped = el.value.replace(prefix, "", 1).lstrip("-\n")
+                    return sub_cls(stripped, el.title)
+
+        return el
+
+
+class DocstringSectionSeeAlso(_DocstringSectionPatched):
+    kind = DocstringSectionKindPatched.see_also
+
+
+class DocstringSectionNotes(_DocstringSectionPatched):
+    kind = DocstringSectionKindPatched.notes
+
+
+class DocstringSectionWarnings(_DocstringSectionPatched):
+    kind = DocstringSectionKindPatched.warnings
+
+
+# Patch Example elements ------------------------------------------------------
+
+
+@dataclass
+class ExampleCode:
+    value: str
+
+
+@dataclass
+class ExampleText:
+    value: str
+
+
+def tuple_to_data(el: "tuple[ds.DocstringSectionKind, str]"):
+    """Re-format funky tuple setup in example section to be a class."""
+    assert len(el) == 2
+
+    kind, value = el
+    if kind.value == "examples":
+        return ExampleCode(value)
+    elif kind.value == "text":
+        return ExampleText(value)
+
+    raise ValueError(f"Unsupported first element in tuple: {kind}")
 
 
 # Tree previewer ==============================================================
@@ -106,7 +203,7 @@ def __init__(self, string_max_length: int = 50, max_depth=999):
     def format(self, call, depth=0, pad=0):
         """Return a Symbolic or Call back as a nice tree, with boxes for nodes."""
 
-        call = self.transform(call)
+        call = transform(call)
 
         crnt_fields = fields(call)
 
@@ -147,19 +244,6 @@ def get_field(self, obj, k):
 
         return getattr(obj, k)
 
-    def transform(self, obj):
-        # TODO: currently this transform happens here, and in the renderer.
-        #       let's consolidate this into one step (when getting the object)
-        if isinstance(obj, tuple):
-            try:
-                return tuple_to_data(obj)
-            except ValueError:
-                pass
-        elif isinstance(obj, ds.DocstringSectionText):
-            return docstring_section_narrow(obj)
-
-        return obj
-
     def fmt_pipe(self, x, is_final=False, pad=0):
         if not is_final:
             connector = self.icon_connector if not is_final else "  "

quartodoc/renderers.py

@@ -1,91 +1,14 @@
+import quartodoc.ast as qast
 import re
 
-from enum import Enum
 from griffe.docstrings import dataclasses as ds
 from griffe import dataclasses as dc
-from dataclasses import dataclass
 from tabulate import tabulate
 from plum import dispatch
 from typing import Tuple, Union
 
 
-# Docstring rendering =========================================================
-
 # utils -----------------------------------------------------------------------
-# these largely re-format the output of griffe
-
-
-def tuple_to_data(el: "tuple[ds.DocstringSectionKind, str]"):
-    """Re-format funky tuple setup in example section to be a class."""
-    assert len(el) == 2
-
-    kind, value = el
-    if kind.value == "examples":
-        return ExampleCode(value)
-    elif kind.value == "text":
-        return ExampleText(value)
-
-    raise ValueError(f"Unsupported first element in tuple: {kind}")
-
-
-def docstring_section_narrow(el: ds.DocstringSection) -> ds.DocstringSection:
-    # attempt to narrow down text sections
-    to_narrow = [
-        DocstringSectionSeeAlso,
-        DocstringSectionNotes,
-        DocstringSectionWarnings,
-    ]
-    prefix = "See Also\n---"
-
-    if isinstance(el, ds.DocstringSectionText):
-        for cls in to_narrow:
-            prefix = cls.kind.value.title() + "\n---"
-            if el.value.lstrip("\n").startswith(prefix):
-
-                stripped = el.value.replace(prefix, "", 1).lstrip("-\n")
-                return cls(stripped, el.title)
-
-    return el
-
-
-class DocstringSectionKindPatched(Enum):
-    see_also = "see also"
-    notes = "notes"
-    warnings = "warnings"
-
-
-class DocstringSectionSeeAlso(ds.DocstringSection):
-    kind = DocstringSectionKindPatched.see_also
-
-    def __init__(self, value: str, title: "str | None"):
-        self.value = value
-        super().__init__(title)
-
-
-class DocstringSectionNotes(ds.DocstringSection):
-    kind = DocstringSectionKindPatched.notes
-
-    def __init__(self, value: str, title: "str | None"):
-        self.value = value
-        super().__init__(title)
-
-
-class DocstringSectionWarnings(ds.DocstringSection):
-    kind = DocstringSectionKindPatched.warnings
-
-    def __init__(self, value: str, title: "str | None"):
-        self.value = value
-        super().__init__(title)
-
-
-@dataclass
-class ExampleCode:
-    value: str
-
-
-@dataclass
-class ExampleText:
-    value: str
 
 
 def escape(val: str):
@@ -258,7 +181,7 @@ def render(self, el: Union[dc.Object, dc.Alias]):
             pass
         else:
             for section in el.docstring.parsed:
-                new_el = docstring_section_narrow(section)
+                new_el = qast.transform(section)
                 title = new_el.kind.value
                 body = self.render(new_el)
 
@@ -309,7 +232,7 @@ def render(self, el: dc.Parameter):
     # or a section with a header not included in the numpydoc standard
     @dispatch
     def render(self, el: ds.DocstringSectionText):
-        new_el = docstring_section_narrow(el)
+        new_el = qast.transform(el)
         if isinstance(new_el, ds.DocstringSectionText):
             # ensures we don't recurse forever
             return el.value
@@ -349,7 +272,7 @@ def render(self, el: ds.DocstringAttribute):
     # see also ----
 
     @dispatch
-    def render(self, el: DocstringSectionSeeAlso):
+    def render(self, el: qast.DocstringSectionSeeAlso):
         # TODO: attempt to parse See Also sections
         return convert_rst_link_to_md(el.value)
 
@@ -358,11 +281,11 @@ def render(self, el: DocstringSectionSeeAlso):
     @dispatch
     def render(self, el: ds.DocstringSectionExamples):
         # its value is a tuple: DocstringSectionKind["text" | "examples"], str
-        data = map(tuple_to_data, el.value)
+        data = map(qast.transform, el.value)
         return "\n\n".join(list(map(self.render, data)))
 
     @dispatch
-    def render(self, el: ExampleCode):
+    def render(self, el: qast.ExampleCode):
         return f"""```python
 {el.value}
 ```"""
@@ -384,7 +307,7 @@ def render(self, el: Union[ds.DocstringReturn, ds.DocstringRaise]):
     # unsupported parts ----
 
     @dispatch
-    def render(self, el: ExampleText):
+    def render(self, el: qast.ExampleText):
         return el.value
 
     @dispatch.multi(

quartodoc/tests/test_ast.py

@@ -0,0 +1,61 @@
+import quartodoc.ast as qast
+
+import pytest
+from griffe.docstrings import dataclasses as ds
+from griffe import dataclasses as dc
+from griffe.docstrings.parsers import parse_numpy
+
+
+@pytest.mark.parametrize(
+    "el, body",
+    [
+        ("See Also\n---", ""),
+        ("See Also\n---\n", ""),
+        ("See Also\n--------", ""),
+        ("\n\nSee Also\n---\n", ""),
+        ("See Also\n---\nbody text", "body text"),
+        ("See Also\n---\nbody text", "body text"),
+    ],
+)
+def test_transform_docstring_section(el, body):
+    src = ds.DocstringSectionText(el, title=None)
+    res = qast._DocstringSectionPatched.transform(src)
+
+    assert isinstance(res, qast.DocstringSectionSeeAlso)
+    assert res.value == body
+
+
+@pytest.mark.parametrize(
+    "el, cls",
+    [
+        ("See Also\n---", qast.DocstringSectionSeeAlso),
+        ("Warnings\n---", qast.DocstringSectionWarnings),
+        ("Notes\n---", qast.DocstringSectionNotes),
+    ],
+)
+def test_transform_docstring_section_subtype(el, cls):
+    # using transform method ----
+    src = ds.DocstringSectionText(el, title=None)
+    res = qast._DocstringSectionPatched.transform(src)
+
+    assert isinstance(res, cls)
+
+    # using transform function ----
+    parsed = parse_numpy(dc.Docstring(el))
+    assert len(parsed) == 1
+
+    res2 = qast.transform(parsed[0])
+    assert isinstance(res2, cls)
+
+
+@pytest.mark.xfail(reason="TODO: sections get grouped into single element")
+def test_transform_docstring_section_clump():
+    docstring = "See Also---\n\nWarnings\n---\n\nNotes---\n\n"
+    parsed = parse_numpy(dc.Docstring(docstring))
+
+    assert len(parsed) == 1
+
+    # res = transform(parsed[0])
+
+    # what to do here? this should more reasonably be handled when transform
+    # operates on the root.