Fix async iterable not showing up attributetable

Co-authored-by: Danny <[email protected]>

docs/extensions/attributetable.py

@@ -1,78 +1,105 @@
-from sphinx.util.docutils import SphinxDirective
-from sphinx.locale import _
-from docutils import nodes
-from sphinx import addnodes
-
-from collections import OrderedDict, namedtuple
+from __future__ import annotations
 import importlib
 import inspect
-import os
 import re
+from typing import Dict, List, NamedTuple, Optional, Tuple, Sequence, TYPE_CHECKING
+
+from docutils import nodes
+from sphinx import addnodes
+from sphinx.application import Sphinx
+from sphinx.environment import BuildEnvironment
+from sphinx.locale import _
+from sphinx.util.docutils import SphinxDirective
+from sphinx.util.typing import OptionSpec
+
+if TYPE_CHECKING:
+    from .builder import DPYHTML5Translator
+
 
 class attributetable(nodes.General, nodes.Element):
     pass
 
+
 class attributetablecolumn(nodes.General, nodes.Element):
     pass
 
+
 class attributetabletitle(nodes.TextElement):
     pass
 
+
 class attributetableplaceholder(nodes.General, nodes.Element):
     pass
 
+
 class attributetablebadge(nodes.TextElement):
     pass
 
+
 class attributetable_item(nodes.Part, nodes.Element):
     pass
 
-def visit_attributetable_node(self, node):
-    class_ = node["python-class"]
+
+def visit_attributetable_node(self: DPYHTML5Translator, node: attributetable) -> None:
+    class_ = node['python-class']
     self.body.append(f'<div class="py-attribute-table" data-move-to-id="{class_}">')
 
-def visit_attributetablecolumn_node(self, node):
+
+def visit_attributetablecolumn_node(self: DPYHTML5Translator, node: attributetablecolumn) -> None:
     self.body.append(self.starttag(node, 'div', CLASS='py-attribute-table-column'))
 
-def visit_attributetabletitle_node(self, node):
+
+def visit_attributetabletitle_node(self: DPYHTML5Translator, node: attributetabletitle) -> None:
     self.body.append(self.starttag(node, 'span'))
 
-def visit_attributetablebadge_node(self, node):
+
+def visit_attributetablebadge_node(self: DPYHTML5Translator, node: attributetablebadge) -> None:
     attributes = {
         'class': 'py-attribute-table-badge',
         'title': node['badge-type'],
     }
     self.body.append(self.starttag(node, 'span', **attributes))
 
-def visit_attributetable_item_node(self, node):
+
+def visit_attributetable_item_node(self: DPYHTML5Translator, node: attributetable_item) -> None:
     self.body.append(self.starttag(node, 'li', CLASS='py-attribute-table-entry'))
 
-def depart_attributetable_node(self, node):
+
+def depart_attributetable_node(self: DPYHTML5Translator, node: attributetable) -> None:
     self.body.append('</div>')
 
-def depart_attributetablecolumn_node(self, node):
+
+def depart_attributetablecolumn_node(self: DPYHTML5Translator, node: attributetablecolumn) -> None:
     self.body.append('</div>')
 
-def depart_attributetabletitle_node(self, node):
+
+def depart_attributetabletitle_node(self: DPYHTML5Translator, node: attributetabletitle) -> None:
     self.body.append('</span>')
 
-def depart_attributetablebadge_node(self, node):
+
+def depart_attributetablebadge_node(self: DPYHTML5Translator, node: attributetablebadge) -> None:
     self.body.append('</span>')
 
-def depart_attributetable_item_node(self, node):
+
+def depart_attributetable_item_node(self: DPYHTML5Translator, node: attributetable_item) -> None:
     self.body.append('</li>')
 
+
 _name_parser_regex = re.compile(r'(?P<module>[\w.]+\.)?(?P<name>\w+)')
 
+
 class PyAttributeTable(SphinxDirective):
     has_content = False
     required_arguments = 1
     optional_arguments = 0
     final_argument_whitespace = False
-    option_spec = {}
+    option_spec: OptionSpec = {}
 
-    def parse_name(self, content):
-        path, name = _name_parser_regex.match(content).groups()
+    def parse_name(self, content: str) -> Tuple[str, str]:
+        match = _name_parser_regex.match(content)
+        if match is None:
+            raise RuntimeError(f"content {content} somehow doesn't match regex in {self.env.docname}.")
+        path, name = match.groups()
         if path:
             modulename = path.rstrip('.')
         else:
@@ -80,11 +107,11 @@ class PyAttributeTable(SphinxDirective):
             if not modulename:
                 modulename = self.env.ref_context.get('py:module')
         if modulename is None:
-            raise RuntimeError('modulename somehow None for %s in %s.' % (content, self.env.docname))
+            raise RuntimeError(f'modulename somehow None for {content} in {self.env.docname}.')
 
         return modulename, name
 
-    def run(self):
+    def run(self) -> List[attributetableplaceholder]:
         """If you're curious on the HTML this is meant to generate:
 
         <div class="py-attribute-table">
@@ -120,17 +147,21 @@ class PyAttributeTable(SphinxDirective):
         node['python-full-name'] = f'{modulename}.{name}'
         return [node]
 
-def build_lookup_table(env):
+
+def build_lookup_table(env: BuildEnvironment) -> Dict[str, List[str]]:
     # Given an environment, load up a lookup table of
     # full-class-name: objects
     result = {}
     domain = env.domains['py']
 
     ignored = {
-        'data', 'exception', 'module', 'class',
+        'data',
+        'exception',
+        'module',
+        'class',
     }
 
-    for (fullname, _, objtype, docname, _, _) in domain.get_objects():
+    for fullname, _, objtype, docname, _, _ in domain.get_objects():
         if objtype in ignored:
             continue
 
@@ -143,9 +174,13 @@ def build_lookup_table(env):
     return result
 
 
-TableElement = namedtuple('TableElement', 'fullname label badge')
+class TableElement(NamedTuple):
+    fullname: str
+    label: str
+    badge: Optional[attributetablebadge]
 
-def process_attributetable(app, doctree, fromdocname):
+
+def process_attributetable(app: Sphinx, doctree: nodes.Node, fromdocname: str) -> None:
     env = app.builder.env
 
     lookup = build_lookup_table(env)
@@ -165,14 +200,17 @@ def process_attributetable(app, doctree, fromdocname):
         else:
             node.replace_self([table])
 
-def get_class_results(lookup, modulename, name, fullname):
+
+def get_class_results(
+    lookup: Dict[str, List[str]], modulename: str, name: str, fullname: str
+) -> Dict[str, List[TableElement]]:
     module = importlib.import_module(modulename)
     cls = getattr(module, name)
 
-    groups = OrderedDict([
-        (_('Attributes'), []),
-        (_('Methods'), []),
-    ])
+    groups: Dict[str, List[TableElement]] = {
+        _('Attributes'): [],
+        _('Methods'): [],
+    }
 
     try:
         members = lookup[fullname]
@@ -205,9 +243,13 @@ def get_class_results(lookup, modulename, name, fullname):
             elif inspect.isfunction(value):
                 if doc.startswith(('A decorator', 'A shortcut decorator')):
                     # finicky but surprisingly consistent
+                    key = _('Methods')
                     badge = attributetablebadge('@', '@')
                     badge['badge-type'] = _('decorator')
+                elif inspect.isasyncgenfunction(value):
                     key = _('Methods')
+                    badge = attributetablebadge('async for', 'async for')
+                    badge['badge-type'] = _('async iterable')
                 else:
                     key = _('Methods')
                     badge = attributetablebadge('def', 'def')
@@ -217,14 +259,14 @@ def get_class_results(lookup, modulename, name, fullname):
 
     return groups
 
-def class_results_to_node(key, elements):
+
+def class_results_to_node(key: str, elements: Sequence[TableElement]) -> attributetablecolumn:
     title = attributetabletitle(key, key)
     ul = nodes.bullet_list('')
     for element in elements:
-        ref = nodes.reference('', '', internal=True,
-                                      refuri='#' + element.fullname,
-                                      anchorname='',
-                                      *[nodes.Text(element.label)])
+        ref = nodes.reference(
+            '', '', internal=True, refuri=f'#{element.fullname}', anchorname='', *[nodes.Text(element.label)]
+        )
         para = addnodes.compact_paragraph('', '', ref)
         if element.badge is not None:
             ul.append(attributetable_item('', element.badge, para))
@@ -233,7 +275,8 @@ def class_results_to_node(key, elements):
 
     return attributetablecolumn('', title, ul)
 
-def setup(app):
+
+def setup(app: Sphinx) -> None:
     app.add_directive('attributetable', PyAttributeTable)
     app.add_node(attributetable, html=(visit_attributetable_node, depart_attributetable_node))
     app.add_node(attributetablecolumn, html=(visit_attributetablecolumn_node, depart_attributetablecolumn_node))