|
| 1 | +from enum import Enum |
| 2 | +from dataclasses import dataclass |
1 | 3 | from griffe.docstrings import dataclasses as ds
|
2 | 4 | from griffe import dataclasses as dc
|
3 | 5 | from plum import dispatch
|
4 | 6 | from typing import Union
|
5 | 7 |
|
6 |
| -# TODO: these classes are created to wrap some tuple outputs |
7 |
| -# we should consolidate logic for transforming the griffe |
8 |
| -# docstring here (or open a griffe issue). |
9 |
| -from .renderers import tuple_to_data, docstring_section_narrow, ExampleCode, ExampleText |
| 8 | + |
| 9 | +# Transform and patched-in classes ============================================ |
| 10 | +# TODO: annotate transform return types. make sure subtypes inherit from correct |
| 11 | +# griffe base objects. |
| 12 | +# TODO: it seems like transform should happen on the root, not individual elements. |
| 13 | + |
| 14 | + |
| 15 | +def transform(el): |
| 16 | + """Return a more specific docstring element, or simply return the original one.""" |
| 17 | + |
| 18 | + if isinstance(el, tuple): |
| 19 | + try: |
| 20 | + return tuple_to_data(el) |
| 21 | + except ValueError: |
| 22 | + pass |
| 23 | + elif isinstance(el, ds.DocstringSection): |
| 24 | + return _DocstringSectionPatched.transform(el) |
| 25 | + |
| 26 | + return el |
| 27 | + |
| 28 | + |
| 29 | +# Patch DocstringSection ------------------------------------------------------ |
| 30 | + |
| 31 | + |
| 32 | +class DocstringSectionKindPatched(Enum): |
| 33 | + see_also = "see also" |
| 34 | + notes = "notes" |
| 35 | + warnings = "warnings" |
| 36 | + |
| 37 | + |
| 38 | +class _DocstringSectionPatched(ds.DocstringSection): |
| 39 | + _registry: "dict[Enum, _DocstringSectionPatched]" = {} |
| 40 | + |
| 41 | + def __init__(self, value: str, title: "str | None" = None): |
| 42 | + self.value = value |
| 43 | + super().__init__(title) |
| 44 | + |
| 45 | + def __init_subclass__(cls, **kwargs): |
| 46 | + super().__init_subclass__(**kwargs) |
| 47 | + |
| 48 | + if cls.kind.value in cls._registry: |
| 49 | + raise KeyError(f"A section for kind {cls.kind} already exists") |
| 50 | + |
| 51 | + cls._registry[cls.kind] = cls |
| 52 | + |
| 53 | + @classmethod |
| 54 | + def transform(cls, el: ds.DocstringSection) -> ds.DocstringSection: |
| 55 | + """Attempt to cast DocstringSection element to more specific section type. |
| 56 | +
|
| 57 | + Note that this is meant to patch cases where the general DocstringSectionText |
| 58 | + class represents a section like See Also, etc.. |
| 59 | + """ |
| 60 | + |
| 61 | + if isinstance(el, ds.DocstringSectionText): |
| 62 | + for kind, sub_cls in cls._registry.items(): |
| 63 | + prefix = kind.value.title() + "\n---" |
| 64 | + if el.value.lstrip("\n").startswith(prefix): |
| 65 | + stripped = el.value.replace(prefix, "", 1).lstrip("-\n") |
| 66 | + return sub_cls(stripped, el.title) |
| 67 | + |
| 68 | + return el |
| 69 | + |
| 70 | + |
| 71 | +class DocstringSectionSeeAlso(_DocstringSectionPatched): |
| 72 | + kind = DocstringSectionKindPatched.see_also |
| 73 | + |
| 74 | + |
| 75 | +class DocstringSectionNotes(_DocstringSectionPatched): |
| 76 | + kind = DocstringSectionKindPatched.notes |
| 77 | + |
| 78 | + |
| 79 | +class DocstringSectionWarnings(_DocstringSectionPatched): |
| 80 | + kind = DocstringSectionKindPatched.warnings |
| 81 | + |
| 82 | + |
| 83 | +# Patch Example elements ------------------------------------------------------ |
| 84 | + |
| 85 | + |
| 86 | +@dataclass |
| 87 | +class ExampleCode: |
| 88 | + value: str |
| 89 | + |
| 90 | + |
| 91 | +@dataclass |
| 92 | +class ExampleText: |
| 93 | + value: str |
| 94 | + |
| 95 | + |
| 96 | +def tuple_to_data(el: "tuple[ds.DocstringSectionKind, str]"): |
| 97 | + """Re-format funky tuple setup in example section to be a class.""" |
| 98 | + assert len(el) == 2 |
| 99 | + |
| 100 | + kind, value = el |
| 101 | + if kind.value == "examples": |
| 102 | + return ExampleCode(value) |
| 103 | + elif kind.value == "text": |
| 104 | + return ExampleText(value) |
| 105 | + |
| 106 | + raise ValueError(f"Unsupported first element in tuple: {kind}") |
10 | 107 |
|
11 | 108 |
|
12 | 109 | # Tree previewer ==============================================================
|
@@ -106,7 +203,7 @@ def __init__(self, string_max_length: int = 50, max_depth=999):
|
106 | 203 | def format(self, call, depth=0, pad=0):
|
107 | 204 | """Return a Symbolic or Call back as a nice tree, with boxes for nodes."""
|
108 | 205 |
|
109 |
| - call = self.transform(call) |
| 206 | + call = transform(call) |
110 | 207 |
|
111 | 208 | crnt_fields = fields(call)
|
112 | 209 |
|
@@ -147,19 +244,6 @@ def get_field(self, obj, k):
|
147 | 244 |
|
148 | 245 | return getattr(obj, k)
|
149 | 246 |
|
150 |
| - def transform(self, obj): |
151 |
| - # TODO: currently this transform happens here, and in the renderer. |
152 |
| - # let's consolidate this into one step (when getting the object) |
153 |
| - if isinstance(obj, tuple): |
154 |
| - try: |
155 |
| - return tuple_to_data(obj) |
156 |
| - except ValueError: |
157 |
| - pass |
158 |
| - elif isinstance(obj, ds.DocstringSectionText): |
159 |
| - return docstring_section_narrow(obj) |
160 |
| - |
161 |
| - return obj |
162 |
| - |
163 | 247 | def fmt_pipe(self, x, is_final=False, pad=0):
|
164 | 248 | if not is_final:
|
165 | 249 | connector = self.icon_connector if not is_final else " "
|
|
0 commit comments