Patch for NumpyDocString to add more parameter supported classes

fastcore/docscrape.py

@@ -25,6 +25,7 @@
 from warnings import warn
 from collections import namedtuple
 from collections.abc import Mapping
+from typing import Any
 
 __all__ = ['Parameter', 'NumpyDocString', 'dedent_lines']
 
@@ -93,23 +94,78 @@ def __str__(self):
 
 
 SECTIONS = 'Summary Extended Yields Receives Other Raises Warns Warnings See Also Notes References Examples Attributes Methods'.split()
+'''Named `numpydoc` sections see: https://numpydoc.readthedocs.io/en/latest/format.html#sections'''
+
+PARAM_SECTIONS = {
+    "Parameters",
+    "Other Parameters",
+    "Attributes",
+    "Methods",
+    "Raises",
+    "Warns",
+    "Yields",
+    "Receives"
+}
+'''Set of `numpydoc` sections which should support parameters via `Parameter`.'''
+
 
 class NumpyDocString(Mapping):
     "Parses a numpydoc string to an abstract representation"
     # See the NumPy Doc Manual https://numpydoc.readthedocs.io/en/latest/format.html>
+    # TODO: flushout docstring
+
+    # NOTE: unclear why these are class variables
     sections = {o:[] for o in SECTIONS}
     sections['Summary'] = ['']
+
+    # NOTE: unclear why these are not included in `SECTIONS` given initialization above creates lists
     sections['Parameters'] = []
     sections['Returns'] = []
 
-    def __init__(self, docstring, config=None):
+    # NOTE: following above style, adding `param_sections` as class variable
+    param_sections: set[str] = set(PARAM_SECTIONS)
+
+    def __init__(
+        self, docstring, 
+        config=None, # TODO: figure this out
+        supported_sections: list[str] | None = SECTIONS,
+        supports_params: set[str] | None = PARAM_SECTIONS
+    ):
+
+        # If None, set to default supported set
+        if supports_params is None: supports_params = set(PARAM_SECTIONS)
+        else:
+            # add missing to class variable
+            missing = set(supports_params) - set(self.param_sections)
+            for sec in missing: self.param_sections.add(sec)
+
+        # If None, set to default supported set
+        if supported_sections is None: supported_sections = set(SECTIONS)
+        else:
+            # add missing to class variable
+            missing = set(supported_sections) - set(self.sections.keys())
+            for sec in missing: self.sections[sec] = []
+
+
+        # --- original initialization ---
         docstring = textwrap.dedent(docstring).split('\n')
         self._doc = Reader(docstring)
         self._parsed_data = copy.deepcopy(self.sections)
         self._parse()
+
+        # --- fastcore default normalization ---
         self['Parameters'] = {o.name:o for o in self['Parameters']}
         if self['Returns']: self['Returns'] = self['Returns'][0]
-        for section in SECTIONS: self[section] = dedent_lines(self[section], split=False)
+
+        # --- our patch: normalize ALL parameter-like sections ---
+        for sec in supports_params:
+            if sec in self._parsed_data:
+                self._parsed_data[sec] = self._normalize_param_section(self._parsed_data[sec])
+
+
+        # --- continue normal fastcore behavior ---
+        for section in SECTIONS: 
+            self[section] = dedent_lines(self[section], split=False)
 
     def __iter__(self): return iter(self._parsed_data)
     def __len__(self): return len(self._parsed_data)
@@ -174,6 +230,25 @@ def _parse_param_list(self, content, single_element_is_type=False):
             params.append(Parameter(arg_name, arg_type, desc))
         return params
 
+    def _normalize_param_section(self, val: list[Parameter] | Any) -> dict[Parameter] | Any:
+        """
+        Convert lists of `Parameter` objects into a dict or clean list.
+        """
+        # Not a list? Then noop.
+        if not isinstance(val, list):
+            return val
+
+        # Falsy value i.e. empty list? Then noop.
+        if not val:
+            return val
+
+        # Lazy check, assumes if first value is a Parameter, all are.
+        if not isinstance(val[0], Parameter):
+            return val
+
+        # Convert to dict[name -> Parameter]
+        return {p.name: p for p in val}
+
     def _parse_summary(self):
         """Grab signature (if given) and summary"""
         if self._is_at_section(): return