doc-forge/docforge/loaders/griffe_loader.py

"""
Utilities for loading and introspecting Python modules using Griffe.

This module provides the ``GriffeLoader`` class and helper utilities used to
discover Python modules, introspect their structure, and convert the results
into doc-forge documentation models.
"""

import logging
from pathlib import Path
from typing import List, Optional

from griffe import (
    GriffeLoader as _GriffeLoader,
    ModulesCollection,
    LinesCollection,
    Object,
    AliasResolutionError,
)

from docforge.models import Module, Project, DocObject

logger = logging.getLogger(__name__)


def discover_module_paths(
    module_name: str,
    project_root: Path | None = None,
) -> List[str]:
    """
    Discover Python modules within a package directory.

    The function scans the filesystem for ``.py`` files inside the specified
    package and converts them into dotted module import paths.

    Discovery rules:

    - Directories containing ``__init__.py`` are treated as packages.
    - Each ``.py`` file is treated as a module.
    - Results are returned as dotted import paths.

    Args:
        module_name: Top-level package name to discover modules from.
        project_root: Root directory used to resolve module paths. If not
            provided, the current working directory is used.

    Returns:
        A sorted list of unique dotted module import paths.

    Raises:
        FileNotFoundError: If the specified package directory does not exist.
    """

    if project_root is None:
        project_root = Path.cwd()

    pkg_dir = project_root / module_name
    if not pkg_dir.exists():
        raise FileNotFoundError(f"Package not found: {pkg_dir}")

    module_paths: List[str] = []

    for path in pkg_dir.rglob("*.py"):
        if path.name == "__init__.py":
            module_path = path.parent
        else:
            module_path = path

        rel = module_path.relative_to(project_root)
        dotted = ".".join(rel.with_suffix("").parts)
        module_paths.append(dotted)

    return sorted(set(module_paths))


class GriffeLoader:
    """
    Load Python modules using Griffe and convert them into doc-forge models.

    This loader uses the Griffe introspection engine to analyze Python source
    code and transform the extracted information into ``Project``, ``Module``,
    and ``DocObject`` instances used by doc-forge.
    """

    def __init__(self) -> None:
        """
        Initialize the Griffe-backed loader.

        Creates an internal Griffe loader instance with dedicated collections
        for modules and source lines.
        """
        self._loader = _GriffeLoader(
            modules_collection=ModulesCollection(),
            lines_collection=LinesCollection(),
        )

    def load_project(
            self,
            module_paths: List[str],
            project_name: Optional[str] = None,
            skip_import_errors: bool = None,
    ) -> Project:
        """
        Load multiple modules and assemble them into a Project model.

        Each module path is introspected and converted into a ``Module``
        instance. All modules are then aggregated into a single ``Project``
        object.

        Args:
            module_paths: List of dotted module import paths to load.
            project_name: Optional override for the project name. Defaults
                to the top-level name of the first module.
            skip_import_errors: If True, modules that fail to load will be
                skipped instead of raising an error.

        Returns:
            A populated ``Project`` instance containing the loaded modules.

        Raises:
            ValueError: If no module paths are provided.
            ImportError: If a module fails to load and
                ``skip_import_errors`` is False.
        """
        if not module_paths:
            raise ValueError("At least one module path must be provided")

        if project_name is None:
            project_name = module_paths[0].split(".")[0]

        project = Project(name=project_name)

        for module_path in module_paths:
            try:
                module = self.load_module(module_path)
            except ImportError as import_error:
                if skip_import_errors:
                    logger.debug("Could not load %s: %s", module_path, import_error)
                    continue
                else:
                    raise import_error
            project.add_module(module)

        return project

    def load_module(self, path: str) -> Module:
        """
        Load and convert a single Python module.

        The module is introspected using Griffe and then transformed into
        a doc-forge ``Module`` model.

        Args:
            path: Dotted import path of the module.

        Returns:
            A populated ``Module`` instance.
        """
        self._loader.load(path)
        griffe_module = self._loader.modules_collection[path]

        return self._convert_module(griffe_module)

    # -------------------------
    # Conversion helpers
    # -------------------------

    def _convert_module(self, obj: Object) -> Module:
        """
        Convert a Griffe module object into a doc-forge Module.

        All public members of the module are recursively converted into
        ``DocObject`` instances.

        Args:
            obj: Griffe object representing the module.

        Returns:
            A populated ``Module`` model.
        """
        module = Module(
            path=obj.path,
            docstring=self._safe_docstring(obj),
        )

        for name, member in obj.members.items():
            if name.startswith("_"):
                continue

            module.add_object(self._convert_object(member))

        return module

    def _convert_object(self, obj: Object) -> DocObject:
        """
        Convert a Griffe object into a doc-forge DocObject.

        The conversion preserves the object's metadata such as name,
        kind, path, signature, and docstring. Child members are processed
        recursively.

        Args:
            obj: Griffe object representing a documented Python object.

        Returns:
            A ``DocObject`` instance representing the converted object.
        """
        kind = obj.kind.value
        signature = self._safe_signature(obj)

        doc_obj = DocObject(
            name=obj.name,
            kind=kind,
            path=obj.path,
            signature=signature,
            docstring=self._safe_docstring(obj),
        )

        try:
            for name, member in obj.members.items():
                if name.startswith("_"):
                    continue
                doc_obj.add_member(self._convert_object(member))
        except AliasResolutionError:
            pass

        return doc_obj

    # -------------------------
    # Safe extractors
    # -------------------------

    def _safe_docstring(self, obj: Object) -> Optional[str]:
        """
        Safely extract a docstring from a Griffe object.

        Args:
            obj: Griffe object to inspect.

        Returns:
            The raw docstring text if available, otherwise ``None``.
        """
        try:
            return obj.docstring.value if obj.docstring else None
        except AliasResolutionError:
            return None

    def _safe_signature(self, obj: Object) -> Optional[str]:
        """
        Safely extract the signature of a Griffe object.

        Args:
            obj: Griffe object to inspect.

        Returns:
            String representation of the object's signature if available,
            otherwise ``None``.
        """
        try:
            if hasattr(obj, "signature") and obj.signature:
                return str(obj.signature)
        except AliasResolutionError:
            return None
        return None