Add annotations to DoctypeTransformer

Author: lagru

src/docstub/_docstrings.py

@@ -295,35 +295,29 @@ def doctype_to_annotation(self, doctype):
             self._collected_imports = None
             self._unknown_qualnames = None
 
-    def __default__(self, data, children, meta):
-        """Unpack children of rule nodes by default.
-
+    def annotation_with_meta(self, tree):
+        """
         Parameters
         ----------
-        data : lark.Token
-            The rule-token of the current node.
-        children : list[lark.Token]
-            The children of the current node.
-        meta : lark.tree.Meta
-            Meta information for the current node.
+        tree : lark.Tree
 
         Returns
         -------
-        out : lark.Token or list[lark.Token]
-            Either a token or list of tokens.
+        out : str
         """
-        if isinstance(children, list) and len(children) == 1:
-            out = children[0]
-            out.type = data.upper()  # Turn rule into "token"
-        else:
-            out = children
-        return out
-
-    def annotation_with_meta(self, tree):
         out = " | ".join(tree.children)
         return out
 
     def qualname(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.Token
+        """
         children = tree.children
         _qualname = ".".join(children)
 
@@ -346,21 +340,57 @@ def qualname(self, tree):
         return _qualname
 
     def rst_role(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.Token
+        """
         qualname = _find_one_token(tree, name="QUALNAME")
         return qualname
 
     def or_expression(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
         out = " | ".join(tree.children)
         return out
 
     def subscription_expression(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
         _container, *_content = tree.children
         _content = ", ".join(_content)
         assert _content
         out = f"{_container}[{_content}]"
         return out
 
     def literal_expression(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
         out = ", ".join(tree.children)
         out = f"Literal[{out}]"
         if self.types_db is not None:
@@ -370,13 +400,31 @@ def literal_expression(self, tree):
         return out
 
     def array_expression(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : str
+        """
         name = _find_one_token(tree, name="ARRAY_NAME")
         children = [child for child in tree.children if child != name]
         if children:
             name = f"{name}[{', '.join(children)}]"
-        return name
+        return str(name)
 
     def array_name(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.Token
+        """
         # Treat `array_name` as `qualname`, but mark it as an array name,
         # so we know which one to treat as the container in `array_expression`
         # This currently relies on a hack that only allows specific names
@@ -386,17 +434,68 @@ def array_name(self, tree):
         return qualname
 
     def shape(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.visitors._DiscardType
+        """
         logger.debug("dropping shape information")
         return lark.Discard
 
     def optional(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.visitors._DiscardType
+        """
         logger.debug("dropping optional / default info")
         return lark.Discard
 
     def extra_info(self, tree):
+        """
+        Parameters
+        ----------
+        tree : lark.Tree
+
+        Returns
+        -------
+        out : lark.visitors._DiscardType
+        """
         logger.debug("dropping extra info")
         return lark.Discard
 
+    def __default__(self, data, children, meta):
+        """Unpack children of rule nodes by default.
+
+        Parameters
+        ----------
+        data : lark.Token
+            The rule-token of the current node.
+        children : list[lark.Token]
+            The children of the current node.
+        meta : lark.tree.Meta
+            Meta information for the current node.
+
+        Returns
+        -------
+        out : lark.Token or list[lark.Token]
+            Either a token or list of tokens.
+        """
+        if isinstance(children, list) and len(children) == 1:
+            out = children[0]
+            out.type = data.upper()  # Turn rule into "token"
+        else:
+            out = children
+        return out
+
     def _match_import(self, qualname, *, meta):
         """Match `qualname` to known imports or alias to "Incomplete".