Simplify generate.py

RustPython · youknowone · Nov 14, 2025 · Nov 1, 2025 · Nov 3, 2025 · Nov 3, 2025
commit c2b2249ae0b7aa8960ed8bbf09c4afb6ff10570d
diff --git a/crates/rustpython_doc_db/generate.py b/crates/rustpython_doc_db/generate.py
@@ -33,26 +33,32 @@
 CRATE_DIR = pathlib.Path(__file__).parent
 OUT_FILE = CRATE_DIR / "src" / f"{sys.platform}.inc.rs"
 
+type Parts = tuple[str, ...]
+
 
 class DocEntry(typing.NamedTuple):
-    parts: tuple[str, ...]
+    parts: Parts
     doc: str | None
 
     @property
     def phf_entry(self) -> str:
-        if self.doc:
-            escaped = re.sub(UNICODE_ESCAPE, r"\\u{\1}", self.doc)
-            dumped = json.dumps(escaped)
-            doc = f"Some({dumped})"
-        else:
-            doc = "None"
+        escaped = re.sub(UNICODE_ESCAPE, r"\\u{\1}", inspect.cleandoc(self.doc))
+        doc = json.dumps(escaped)
 
         key = json.dumps(".".join(self.parts))
         return f"{key} => {doc}"
 
 
 def is_c_extension(module: types.ModuleType) -> bool:
     """
+    Check whether a module was written in C.
+
+    Returns
+    -------
+    bool
+
+    Notes
+    -----
     Adapted from: https://stackoverflow.com/a/39304199
     """
     loader = getattr(module, "__loader__", None)
@@ -73,7 +79,14 @@ def is_c_extension(module: types.ModuleType) -> bool:
     return module_filetype in EXTENSION_SUFFIXES
 
 
-def is_child(obj: typing.Any, module: types.ModuleType) -> bool:
+def is_child_of(obj: typing.Any, module: types.ModuleType) -> bool:
+    """
+    Whether or not an object is a child of a module.
+
+    Returns
+    -------
+    bool
+    """
     return inspect.getmodule(obj) is module
 
 
@@ -82,7 +95,7 @@ def iter_modules() -> "Iterable[types.ModuleType]":
     Yields
     ------
     :class:`types.Module`
-        Modules that are written in C. (not pure python)
+        Python modules.
     """
     for module_name in sys.stdlib_module_names - IGNORED_MODULES:
         try:
@@ -93,20 +106,27 @@ def iter_modules() -> "Iterable[types.ModuleType]":
             warnings.warn(f"Could not import {module_name}", category=ImportWarning)
             continue
 
-        if not is_c_extension(module):
-            continue
-
         yield module
 
 
+def iter_c_modules() -> "Iterable[types.ModuleType]":
+    """
+    Yields
+    ------
+    :class:`types.Module`
+        Modules that are written in C. (not pure python)
+    """
+    yield from filter(is_c_extension, iter_modules())
+
+
 def traverse(
-    obj: typing.Any, module: types.ModuleType, name_parts: tuple[str, ...] = ()
+    obj: typing.Any, module: types.ModuleType, parts: Parts = ()
 ) -> "typing.Iterable[DocEntry]":
     if inspect.ismodule(obj):
-        name_parts += (obj.__name__,)
+        parts += (obj.__name__,)
 
     if any(f(obj) for f in (inspect.ismodule, inspect.isclass, inspect.isbuiltin)):
-        yield DocEntry(name_parts, pydoc._getdoc(obj))
+        yield DocEntry(parts, pydoc._getowndoc(obj))
 
     for name, attr in inspect.getmembers(obj):
         if name in IGNORED_ATTRS:
@@ -115,68 +135,74 @@ def traverse(
         if attr == obj:
             continue
 
-        parts = name_parts + (name,)
-
-        if (module is obj) and (not is_child(attr, module)):
+        if (module is obj) and (not is_child_of(attr, module)):
             continue
 
+        # Don't recurse into modules imported by our module. i.e. `ipaddress.py` imports `re` don't traverse `re`
         if (not inspect.ismodule(obj)) and inspect.ismodule(attr):
             continue
 
-        new_name_parts = name_parts + (name,)
+        new_parts = parts + (name,)
 
         attr_typ = type(attr)
-        is_type_or_module = (attr_typ is type) or (attr_typ is type(__builtins__))
+        is_type_or_builtin = any(attr_typ is x for x in (type, type(__builtins__)))
 
-        if is_type_or_module:
-            yield from traverse(attr, module, new_name_parts)
+        if is_type_or_builtin:
+            yield from traverse(attr, module, new_parts)
             continue
 
-        if (
+        is_callable = (
             callable(attr)
             or not issubclass(attr_typ, type)
             or attr_typ.__name__ in ("getset_descriptor", "member_descriptor")
-        ) or any(
+        )
+
+        is_func = any(
             f(obj)
-            for f in (
-                inspect.isfunction,
-                inspect.ismethod,
-                inspect.ismethoddescriptor,
-            )
-        ):
-            yield DocEntry(new_name_parts, pydoc._getdoc(attr))
+            for f in (inspect.isfunction, inspect.ismethod, inspect.ismethoddescriptor)
+        )
+
+        if is_callable or is_func:
+            yield DocEntry(new_parts, pydoc._getowndoc(attr))
 
 
 def find_doc_entires() -> "Iterable[DocEntry]":
     yield from (
-        doc_entry for module in iter_modules() for doc_entry in traverse(module, module)
+        doc_entry
+        for module in iter_c_modules()
+        for doc_entry in traverse(module, module)
     )
     yield from (doc_entry for doc_entry in traverse(__builtins__, __builtins__))
 
     builtin_types = [
+        type(None),
         type(bytearray().__iter__()),
         type(bytes().__iter__()),
         type(dict().__iter__()),
-        type(dict().values().__iter__()),
+        type(dict().items()),
         type(dict().items().__iter__()),
         type(dict().values()),
-        type(dict().items()),
-        type(set().__iter__()),
+        type(dict().values().__iter__()),
+        type(lambda: ...),
         type(list().__iter__()),
+        type(memoryview(b"").__iter__()),
         type(range(0).__iter__()),
+        type(set().__iter__()),
         type(str().__iter__()),
         type(tuple().__iter__()),
-        type(None),
-        type(lambda: ...),
     ]
     for typ in builtin_types:
-        name_parts = ("builtins", typ.__name__)
-        yield DocEntry(name_parts, pydoc._getdoc(typ))
-        yield from traverse(typ, __builtins__, name_parts)
+        parts = ("builtins", typ.__name__)
+        yield DocEntry(parts, pydoc._getowndoc(typ))
+        yield from traverse(typ, __builtins__, parts)
 
 
 def main():
-    doc_entries = {doc_entry.phf_entry for doc_entry in find_doc_entires()}
+    doc_entries = {
+        doc_entry.phf_entry
+        for doc_entry in find_doc_entires()
+        if doc_entry.doc is not None
+    }
 
     lines = ",\n".join(sorted(doc_entries))
     lines = textwrap.indent(lines, prefix=" " * 4)
@@ -187,9 +213,11 @@ def main():
     out = f"""
 // This file was auto generated by: {script_name}
 // CPython version: {python_version}
-phf::phf_map! {{
+use phf::{{Map, phf_map}};
+
+pub static DB: Map<&'static str, &'static str> = phf_map! {{
 {lines}
-}}
+}};
 """.lstrip()
 
     OUT_FILE.write_text(out)