Simplify part of, type and cleanup py2d

AutoDuck/py2d.py

@@ -1,59 +1,69 @@
+from __future__ import annotations
+
 import re
 import sys
-import types
+from collections.abc import Callable, Iterable
+from functools import partial
+from types import FunctionType, MethodType
+from typing import TYPE_CHECKING, Generator, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from _typeshed import SupportsWrite
+
+_T = TypeVar("_T")
 
 
-def ad_escape(s):
+def ad_escape(s: str) -> str:
     return re.sub(r"([^<]*)<([^>]*)>", r"\g<1>\\<\g<2>\\>", s)
 
 
-Print = __builtins__.__dict__["print"]
+Print: Callable[..., None] = print
 
 
-class DocInfo:
-    def __init__(self, name, ob):
+class DocInfo(Generic[_T]):
+    def __init__(self, name: str, ob: _T) -> None:
+        docstring = (ob.__doc__ or "").strip()
+
+        self.desc = docstring
+        self.short_desc = docstring and docstring.splitlines()[0]
         self.name = name
         self.ob = ob
-        self.short_desc = ""
-        self.desc = ""
 
 
-def BuildArgInfos(ob):
-    ret = []
-    vars = list(ob.__code__.co_varnames[: ob.__code__.co_argcount])
-    vars.reverse()  # for easier default checking.
+class ArgInfo(DocInfo[FunctionType | MethodType]):
+    def __init__(self, name: str, ob: FunctionType | MethodType, default: str) -> None:
+        super().__init__(name, ob)
+        self.desc = name
+        self.short_desc = name
+        self.default = default
+
+
+def BuildArgInfos(ob: FunctionType | MethodType) -> list[ArgInfo]:
+    ret: list[ArgInfo] = []
+    # Reversed for easier default checking.
+    # Since arguments w/ default can only be at the end of a function.
+    vars = reversed(ob.__code__.co_varnames[: ob.__code__.co_argcount])  # type: ignore[union-attr] # false-positive in typeshed
     defs = list(ob.__defaults__ or [])
-    for i, n in enumerate(vars):
-        info = DocInfo(n, ob)
-        info.short_desc = info.desc = n
-        info.default = ""
+    for n in vars:
+        default = ""
         if len(defs):
             default = repr(defs.pop())
-            # the default may be an object, so the repr gives '<...>' - and
-            # the angle brackets screw autoduck.
-            info.default = default.replace("<", "").replace(">", "")
-        ret.append(info)
+            # the default may be an object, so the repr gives '<...>'
+            # and the angle brackets screw AutoDuck.
+            default = default.replace("<", "").replace(">", "")
+        ret.append(ArgInfo(n, ob, default))
     ret.reverse()
     return ret
 
 
-def BuildInfo(name, ob):
-    ret = DocInfo(name, ob)
-    docstring = ob.__doc__ or ""
-    ret.desc = ret.short_desc = docstring.strip()
-    if ret.desc:
-        ret.short_desc = ret.desc.splitlines()[0]
-    return ret
-
+def should_build_function(fn: FunctionType | MethodType) -> bool:
+    return bool(fn.__doc__) and not fn.__name__.startswith("_")
 
-def should_build_function(build_info):
-    return build_info.ob.__doc__ and not build_info.ob.__name__.startswith("_")
 
-
-# docstring aware paragraph generator.  Isn't there something in docutils
-# we can use?
-def gen_paras(val):
-    chunks = []
+# docstring aware paragraph generator.
+# Isn't there something in docutils we can use?
+def gen_paras(val: str) -> Generator[list[str]]:
+    chunks: list[str] = []
     in_docstring = False
     for line in val.splitlines():
         line = ad_escape(line.strip())
@@ -69,19 +79,19 @@ def gen_paras(val):
     yield chunks or [""]
 
 
-def format_desc(desc):
+def format_desc(desc: str) -> str:
     # A little complicated!  Given the docstring for a module, we want to:
     # write:
     # 'first_para_of_docstring'
     # '@comm next para of docstring'
     # '@comm next para of docstring' ... etc
-    # BUT - also handling enbedded doctests, where we write
+    # BUT - also handling embedded doctests, where we write
     # '@iex >>> etc.'
     if not desc:
         return ""
     g = gen_paras(desc)
-    first = next(g)
-    chunks = [first[0]]
+    first = next(g)[0]
+    chunks = [first]
     chunks.extend(["// " + l for l in first[1:]])
     for lines in g:
         first = lines[0]
@@ -94,87 +104,79 @@ def format_desc(desc):
     return "\n".join(chunks)
 
 
-def build_module(fp, mod_name):
+def build_module(mod_name: str) -> None:
     __import__(mod_name)
     mod = sys.modules[mod_name]
-    functions = []
-    classes = []
-    constants = []
+    functions: list[DocInfo[FunctionType]] = []
+    classes: list[DocInfo[type]] = []
+    constants: list[tuple[str, int | str]] = []
     for name, ob in list(mod.__dict__.items()):
         if name.startswith("_"):
             continue
         if hasattr(ob, "__module__") and ob.__module__ != mod_name:
             continue
         if type(ob) == type:
-            classes.append(BuildInfo(name, ob))
-        elif isinstance(ob, types.FunctionType):
-            functions.append(BuildInfo(name, ob))
+            classes.append(DocInfo(name, ob))
+        elif isinstance(ob, FunctionType):
+            if should_build_function(ob):
+                functions.append(DocInfo(name, ob))
         elif name.upper() == name and isinstance(ob, (int, str)):
             constants.append((name, ob))
-    info = BuildInfo(mod_name, mod)
-    Print(f"// @module {mod_name}|{format_desc(info.desc)}", file=fp)
-    functions = [f for f in functions if should_build_function(f)]
+    module_info = DocInfo(mod_name, mod)
+    Print(f"// @module {mod_name}|{format_desc(module_info.desc)}")
     for ob in functions:
-        Print(f"// @pymeth {ob.name}|{ob.short_desc}", file=fp)
+        Print(f"// @pymeth {ob.name}|{ob.short_desc}")
     for ob in classes:
         # only classes with docstrings get printed.
         if not ob.ob.__doc__:
             continue
         ob_name = mod_name + "." + ob.name
-        Print(f"// @pyclass {ob.name}|{ob.short_desc}", file=fp)
+        Print(f"// @pyclass {ob.name}|{ob.short_desc}")
     for ob in functions:
         Print(
             f"// @pymethod |{mod_name}|{ob.name}|{format_desc(ob.desc)}",
-            file=fp,
         )
         for ai in BuildArgInfos(ob.ob):
-            Print(f"// @pyparm |{ai.name}|{ai.default}|{ai.short_desc}", file=fp)
+            Print(f"// @pyparm |{ai.name}|{ai.default}|{ai.short_desc}")
 
     for ob in classes:
         # only classes with docstrings get printed.
         if not ob.ob.__doc__:
             continue
         ob_name = mod_name + "." + ob.name
-        Print(f"// @object {ob_name}|{format_desc(ob.desc)}", file=fp)
-        func_infos = []
+        Print(f"// @object {ob_name}|{format_desc(ob.desc)}")
+        func_infos: list[DocInfo[FunctionType | MethodType]] = []
         # We need to iter the keys then to a getattr() so the funky descriptor
         # things work.
         for n in list(ob.ob.__dict__.keys()):
             o = getattr(ob.ob, n)
-            if isinstance(o, (types.FunctionType, types.MethodType)):
-                info = BuildInfo(n, o)
-                if should_build_function(info):
-                    func_infos.append(info)
+            if isinstance(o, (FunctionType, MethodType)):
+                if should_build_function(o):
+                    func_infos.append(DocInfo(n, o))
         for fi in func_infos:
-            Print(f"// @pymeth {fi.name}|{fi.short_desc}", file=fp)
+            Print(f"// @pymeth {fi.name}|{fi.short_desc}")
         for fi in func_infos:
-            Print(
-                f"// @pymethod |{ob_name}|{fi.name}|{format_desc(fi.desc)}",
-                file=fp,
-            )
+            Print(f"// @pymethod |{ob_name}|{fi.name}|{format_desc(fi.desc)}")
             if hasattr(fi.ob, "im_self") and fi.ob.im_self is ob.ob:
-                Print("// @comm This is a @classmethod.", file=fp)
-            Print(
-                f"// @pymethod |{ob_name}|{fi.name}|{format_desc(fi.desc)}",
-                file=fp,
-            )
+                Print("// @comm This is a @classmethod.")
+            Print(f"// @pymethod |{ob_name}|{fi.name}|{format_desc(fi.desc)}")
             for ai in BuildArgInfos(fi.ob):
-                Print(
-                    f"// @pyparm |{ai.name}|{ai.default}|{ai.short_desc}",
-                    file=fp,
-                )
+                Print(f"// @pyparm |{ai.name}|{ai.default}|{ai.short_desc}")
 
     for name, val in constants:
         desc = f"{name} = {val!r}"
         if isinstance(val, int):
             desc += f" (0x{val:x})"
-        Print(f"// @const {mod_name}|{name}|{desc}", file=fp)
+        Print(f"// @const {mod_name}|{name}|{desc}")
+
 
+def main(fp: SupportsWrite[str], args: Iterable[str]) -> None:
+    global Print
+    Print = partial(print, file=fp)
 
-def main(fp, args):
-    Print("// @doc", file=fp)
+    Print("// @doc")
     for arg in args:
-        build_module(sys.stdout, arg)
+        build_module(arg)
 
 
 if __name__ == "__main__":