Cannot get it to work with native Rust extension

[phoerious]

I have a PyO3 Rust extension for which I'm generating Sphinx documentation. For getting types to show up in the documentation, I have :type: annotations everywhere in my doc strings, which is a pain in the arse.

So I found this Sphinx add-on and thought it could finally replace that. I have a pyi stub files for all modules and an init.py and init.pyi file at the top of my package.

I have installed sphinx_autodoc_typehints via pip and added it to my list of extensions in my conf.py. Unfortunately, when I remove the :type: annotations, the types vanish from the docs.

I tried rebuilding everything from scratch, I moved the extension to various places in the list, but it's not working. Am I missing something?

Thanks!

[gaborbernat]

Do you have the reproducible?

[phoerious]

Yes. This is my module: https://github.com/chatnoir-eu/chatnoir-resiliparse/tree/develop/fastwarc-py/fastwarc

I somewhat got it to work with this patch: chatnoir-eu/chatnoir-resiliparse@a002001#diff-85933aa74a2d66c3e4dcdf7a9ad8397f5a7971080d34ef1108296a7c6b69e7e3R71

def _patch_fastwarc_stub_modules(base_module, submodules, stub_dir):
    """
    Import and patch native extension modules so sphinx_autodoc_typehints can
    find the stubs for them.
    """

    import importlib.util
    from importlib.machinery import SourceFileLoader

    module_imported = importlib.import_module(base_module)

    for name in submodules:
        module_name = f'{base_module}.{name}'
        stub_path = os.path.join(stub_dir, f'{name}.pyi')
        module = getattr(module_imported, name)
        loader = SourceFileLoader(module_name, stub_path)

        module.__file__ = stub_path
        module.__loader__ = loader
        module.__package__ = base_module
        module.__spec__ = importlib.util.spec_from_loader(module_name, loader, origin=stub_path)
        sys.modules[module_name] = module


_patch_fastwarc_stub_modules('fastwarc',
                             ('warc', 'stream_io'),
                             os.path.join(src_dir, 'fastwarc-py', 'fastwarc'))

But it's not resolving the inheritance chains correctly.

Also, this seems to be doing nothing at all? All my Unions are shown as Unions, not pipes:

# Not working??
always_use_bars_union = True

I even managed to vibe-code this monstrosity to make the stub typing work with Sphinx's own stub loading logic (available since 8.2): chatnoir-eu/chatnoir-resiliparse@0a7dafa#diff-85933aa74a2d66c3e4dcdf7a9ad8397f5a7971080d34ef1108296a7c6b69e7e3R89

from importlib.abc import MetaPathFinder
from importlib.machinery import EXTENSION_SUFFIXES, ModuleSpec, SourceFileLoader
import os
import sys

_fastwarc_pkg_dir = os.path.join(src_dir, 'fastwarc-py', 'fastwarc')
_STUBBED_NATIVE_MODULES = {
    'fastwarc.warc': {
        'stub_path': os.path.join(_fastwarc_pkg_dir, 'warc.pyi'),
        'native_mod': lambda: sys.modules['fastwarc']._native.warc,
    },
    'fastwarc.stream_io': {
        'stub_path': os.path.join(_fastwarc_pkg_dir, 'stream_io.pyi'),
        'native_mod': lambda: sys.modules['fastwarc']._native.stream_io,
    },
}

class _NativeStubFinder(MetaPathFinder):
    """
    Make PyO3 submodules look like native extension modules to autodoc.

    Native Sphinx stub loading only activates when ``find_spec()`` returns a
    native-extension origin, for which we need to patch in the .pyi path.
    """

    def __init__(self, module_specs):
        self.module_specs = module_specs

    def find_spec(self, fullname, path=None, target=None):
        spec = self.module_specs.get(fullname)
        if spec is None:
            return None

        mock_native_lib_path = spec['stub_path'].rsplit('.', 1)[0] + EXTENSION_SUFFIXES[0]
        return ModuleSpec(
            fullname,
            SourceFileLoader(fullname, spec['stub_path']),
            origin=mock_native_lib_path
        )


def _copy_docstring(src, dst):
    if not getattr(dst, '__doc__', None) and getattr(src, '__doc__', None):
        dst.__doc__ = src.__doc__


def _hydrate_member_docstrings(dst, src):
    _copy_docstring(src, dst)

    dst_dict = getattr(dst, '__dict__', None)
    src_dict = getattr(src, '__dict__', None)
    if not isinstance(dst_dict, dict) or not isinstance(src_dict, dict):
        return

    for name, dst_member in dst_dict.items():
        src_member = src_dict.get(name)
        if src_member is None:
            continue

        _copy_docstring(src_member, dst_member)

        if isinstance(dst_member, property) and isinstance(src_member, property):
            for accessor_name in ('fget', 'fset', 'fdel'):
                dst_accessor = getattr(dst_member, accessor_name)
                src_accessor = getattr(src_member, accessor_name)
                if dst_accessor is not None and src_accessor is not None:
                    _copy_docstring(src_accessor, dst_accessor)


def setup(_):
    from sphinx.ext.autodoc._dynamic import _importer
    original_importer = _importer._import_module

    sys.meta_path.insert(0, _NativeStubFinder(_STUBBED_NATIVE_MODULES))

    def import_module(modname, try_reload=False):
        # Remove already loaded modules if this is a managed PyO3 module
        if modname in _STUBBED_NATIVE_MODULES:
            for managed_module in _STUBBED_NATIVE_MODULES:
                sys.modules.pop(managed_module, None)

            # Also remove module attribute from parent if set
            parent_name, attr_name = modname.rsplit('.', 1)
            parent_module = sys.modules.get(parent_name)
            if parent_module is not None and hasattr(parent_module, attr_name):
                delattr(parent_module, attr_name)

        # Load again and patch spec
        module = original_importer(modname, try_reload=try_reload)
        if modname in _STUBBED_NATIVE_MODULES:
            # Copy docstrings from native module and its members to loaded stub
            native_module = _STUBBED_NATIVE_MODULES[modname]['native_mod']()
            _copy_docstring(native_module, module)
            for name, member in vars(module).items():
                native_member = getattr(native_module, name, None)
                if native_member is not None:
                    _hydrate_member_docstrings(member, native_member)

        return module

    # Patch _importer._import_module to load stub files properly
    _importer._import_module = import_module

I cleaned it up quite a bit, but it's still a mouthful. Clearly, there must be a better way? On the other hand, this version works quite well. It does resolve all base classes and even shows my Union types with pipes. What's making it very verbose is the necessity to copy over the doc strings if they're not already defined in the stubs.

[phoerious]

Okay, I found a way that works with Sphinx alone and doesn't need this extension. It also doesn't rely anymore on the original submodules being attributes on the native module.

It's still a lot of code, but at least shorter and more robust than the previous attempt. If anyone needs it:

import importlib
from importlib.abc import MetaPathFinder
from importlib.machinery import EXTENSION_SUFFIXES, ModuleSpec, SourceFileLoader
from pathlib import Path
import sys

# Path to your sources root here
src_dir = Path(...)

# Insert your stubbed modules here...
_mymod_pkg_dir = src_dir.joinpath('mymod')
_STUBBED_NATIVE_MODULES = {
    'mymod.submod1': _mymod_pkg_dir.joinpath('submod1.pyi'),
    'mymod.submod2': _mymod_pkg_dir.joinpath('submod2.pyi'),
    # ...
}


class _NativeStubFinder(MetaPathFinder):
    """
    Make PyO3 submodules look like native extension modules to autodoc.

    Native Sphinx stub loading only activates when ``find_spec()`` returns a
    native-extension origin, for which we need to patch in the .pyi path.
    """

    def __init__(self, module_specs):
        self.pyi_paths = module_specs

    def find_spec(self, fullname, path=None, target=None):
        pyi_file: Path = self.pyi_paths.get(fullname)
        if pyi_file is None:
            return None

        mock_native_lib_path = str(pyi_file.parent.joinpath(pyi_file.stem + EXTENSION_SUFFIXES[0]))
        return ModuleSpec(
            fullname,
            SourceFileLoader(fullname, str(pyi_file)),
            origin=mock_native_lib_path
        )


def _copy_docstring(src, dst):
    if dst is not None and not getattr(dst, '__doc__', None) and getattr(src, '__doc__', None):
        dst.__doc__ = src.__doc__


def _hydrate_member_docstrings(dst, src):
    _copy_docstring(src, dst)

    dst_dict = getattr(dst, '__dict__', None)
    src_dict = getattr(src, '__dict__', None)
    if not isinstance(dst_dict, dict) or not isinstance(src_dict, dict):
        return

    for name, dst_member in dst_dict.items():
        src_member = src_dict.get(name)
        if src_member is None:
            continue

        _copy_docstring(src_member, dst_member)

        if isinstance(dst_member, property) and isinstance(src_member, property):
            for accessor_name in ('fget', 'fset', 'fdel'):
                dst_accessor = getattr(dst_member, accessor_name)
                src_accessor = getattr(src_member, accessor_name)
                if dst_accessor is not None and src_accessor is not None:
                    _copy_docstring(src_accessor, dst_accessor)


def setup(_):
    from sphinx.ext.autodoc._dynamic import _importer
    original_importer = _importer._import_module

    native_mods = {}

    for m in _STUBBED_NATIVE_MODULES:
        # Import the parent package once, capture the native submodule object it
        # exposes, then remove the submodule import entry so autodoc can load the
        # stub-backed replacement later.
        parent, name = m.rsplit('.', 1)
        parent_mod = importlib.import_module(parent)
        native_mods[m] = getattr(parent_mod, name)
        sys.modules.pop(m, None)
        if hasattr(parent_mod, name):
            delattr(parent_mod, name)

    sys.meta_path.insert(0, _NativeStubFinder(_STUBBED_NATIVE_MODULES))

    def import_module(modname, try_reload=False):
         # Load new module
        module = original_importer(modname, try_reload=try_reload)
        if modname not in _STUBBED_NATIVE_MODULES:
            return module

        #  Copy docstrings from original module
        _copy_docstring(native_mods[modname], module)
        for name, member in vars(module).items():
            native_member = getattr(native_mods[modname], name, None)
            if native_member is not None:
                _hydrate_member_docstrings(member, native_member)

        return module

    # Patch _importer._import_module to load stub files properly
    _importer._import_module = import_module

[gaborbernat]

This was a bug fixed via 80b518f

[phoerious]

Great, thanks!