Source code for bioimageflow.registry

"""Public tool-registry abstraction.

A :class:`ToolRegistry` wraps :func:`load_versioned_package`,
:func:`resolve_tool_class`, and :func:`serialize_input_schema` /
:func:`serialize_output_schema` behind a single object that GUIs and
other consumers can use without rebuilding metadata serialization or
package-resolution layers themselves.

The registry deliberately separates **install** (a slow, network-bound
side effect) from **register** (a fast, in-process index lookup):

- :meth:`ToolRegistry.install_package` installs a package into the
  tool store, doing the network work.
- :meth:`ToolRegistry.register_package` only loads what is already
  installed and indexes its tools — never installs.

GUIs that validate on every keystroke must call ``register_package``
on hot paths and ``install_package`` only from a user-initiated
action.

Workflow exports may also carry embedded custom tool modules. GUIs can
call :meth:`ToolRegistry.register_workflow` to discover those local tools
for the specific workflow being edited without promoting them to a
versioned tool package.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from bioimageflow.paths import get_tool_store_path
from bioimageflow.tool_loader import (
    ensure_installed,
    load_versioned_package,
)
from bioimageflow.validation import (
    serialize_input_schema,
    serialize_output_schema,
)

logger = logging.getLogger("bioimageflow")
_RegistryKey = tuple[str, str, str, str]




@dataclass(frozen=True)
class ToolMetadata:
    """Public, serialized description of a tool class.

    Produced by :meth:`ToolRegistry.register_package` for every
    BaseTool subclass found in the loaded package.

    Attributes
    ----------
    package
        The tool's package name (e.g. ``"my_tools"``).
    version
        Pinned package version (e.g. ``"1.2.3"``).
    module
        Canonical module path (e.g. ``"my_tools.alpha"``).
    class_name
        The class name as written in the source.
    inputs_schema
        Output of :func:`serialize_input_schema` for this class.
    outputs_schema
        Output of :func:`serialize_output_schema` for this class.
    display_name
        Human-readable label declared on the class, or the class name.
    tags
        Free-form tags declared on the class (empty list if none).
    """
    package: str
    version: str
    module: str
    class_name: str
    inputs_schema: dict[str, Any]
    outputs_schema: dict[str, Any]
    display_name: str
    tags: tuple[str, ...] = field(default_factory=tuple)





class ToolRegistry:
    """Stateful index of tool classes loaded from packages or workflows."""



    def __init__(self, *, store_path: Path | None = None) -> None:
        self._store_path: Path = (
            store_path if store_path is not None else get_tool_store_path()
        )
        self._classes: dict[_RegistryKey, type] = {}
        self._metadata: dict[_RegistryKey, ToolMetadata] = {}


    # -- install vs register ------------------------------------------------



    def install_package(self, name: str, version: str) -> None:
        """Install a versioned package into the tool store.

        This is the slow, network-bound side effect — it does not load
        or index anything. Call :meth:`register_package` separately
        once the install completes.

        ``name`` is the import name (the module a workflow expects to
        import). ``ensure_installed`` infers the PyPI name from this
        by replacing underscores with hyphens, matching the convention
        used by :meth:`Workflow.from_dict`'s auto-installer.
        """
        pypi_name = name.replace("_", "-")
        ensure_installed(name, version, pypi_name, self._store_path)




    def register_package(self, name: str, version: str) -> list[ToolMetadata]:
        """Load an *already installed* package and index its tools.

        Returns the metadata for every BaseTool subclass
        discovered in the package. Raises :class:`FileNotFoundError`
        if the package is not present in the store — the caller must
        install it first via :meth:`install_package` (or skip
        registration on the hot path).
        """
        from bioimageflow_core.tool import BaseTool

        # load_versioned_package raises FileNotFoundError if the package
        # is not installed — we deliberately propagate that error rather
        # than auto-installing.
        mod = load_versioned_package(name, version, self._store_path)

        discovered: list[ToolMetadata] = []
        seen: set[type] = set()
        unstamped: dict[type, str] = {}

        # Walk the loaded module and any submodules registered under its
        # scoped name. Tool classes stamped by _stamp_tool_classes carry
        # _bif_canonical_module pointing at the original module path.
        import sys
        scoped_prefix = mod.__name__
        for sys_name, sys_mod in list(sys.modules.items()):
            if sys_mod is None:
                continue
            if sys_name != scoped_prefix and not sys_name.startswith(
                scoped_prefix + "."
            ):
                continue
            for attr_name in dir(sys_mod):
                try:
                    obj = getattr(sys_mod, attr_name)
                except Exception:
                    continue
                if not isinstance(obj, type):
                    continue
                if obj in seen:
                    continue
                if not issubclass(obj, BaseTool):
                    continue
                if obj is BaseTool:
                    continue
                if getattr(obj, "_bif_package", None) != name:
                    # A tool-like class that wasn't stamped for this
                    # package. If its __module__ points at the canonical
                    # package name, it's almost certainly an absolute
                    # import in __init__.py bypassing the scoped loader.
                    # Track it so we can warn once per offending class
                    # below — silent skip would leave the GUI's tool
                    # list empty with no actionable diagnostic.
                    obj_module = getattr(obj, "__module__", "")
                    if (
                        obj_module == name
                        or obj_module.startswith(name + ".")
                    ) and obj not in unstamped:
                        unstamped[obj] = obj_module
                    continue
                seen.add(obj)

                meta = self._build_metadata(obj)
                key = self._key(meta)
                self._classes[key] = obj
                self._metadata[key] = meta
                discovered.append(meta)

        for cls, obj_module in unstamped.items():
            logger.warning(
                "Tool class %s.%s found in package %r v%s but missing "
                "scope marker — this usually means the package's "
                "__init__.py uses absolute imports "
                "(`from %s.module import X`) instead of relative ones "
                "(`from .module import X`). The class will not be "
                "registered. See specs.md §Tool Packages.",
                obj_module, cls.__name__, name, version, name,
            )

        return discovered




    def register_workflow(self, workflow: Any) -> list[ToolMetadata]:
        """Index custom tools carried by a workflow.

        ``workflow`` may be either a live :class:`bioimageflow.Workflow`
        instance or an exported workflow dict. Live workflows are
        inspected recursively; exported archives carry one deduplicated
        ``custom_sources`` table.
        Package tools referenced by the workflow are not loaded or
        installed here; use :meth:`register_package` for those.
        """
        from bioimageflow.workflow_node import WorkflowNode
        from bioimageflow.workflow import (
            Workflow,
            _is_workflow_custom_class,
            _load_custom_sources,
            _resolve_custom_tool_class,
        )

        classes: list[type] = []
        if isinstance(workflow, Workflow):
            def visit_live(definition: Workflow) -> None:
                for node in definition.nodes.values():
                    if isinstance(node, WorkflowNode):
                        visit_live(node.workflow)
                        continue
                    cls = type(node.tool)
                    if _is_workflow_custom_class(cls):
                        classes.append(cls)

            visit_live(workflow)
        elif isinstance(workflow, dict):
            if set(workflow) == {"archive_version", "workflow", "custom_sources"}:
                modules = _load_custom_sources(workflow["custom_sources"])
                graph = workflow["workflow"]
            else:
                modules = {}
                graph = workflow

            def visit_graph(definition: dict[str, Any]) -> None:
                for node_data in definition.get("nodes", []):
                    if node_data.get("type") == "workflow":
                        visit_graph(node_data["workflow"])
                        continue
                    source_id = node_data.get("source_module")
                    if source_id:
                        classes.append(_resolve_custom_tool_class(
                            modules,
                            source_id,
                            node_data["tool_module"],
                            node_data["tool_class"],
                        ))

            visit_graph(graph)
        else:
            raise TypeError(
                "register_workflow expects a Workflow instance or workflow dict"
            )

        discovered: list[ToolMetadata] = []
        seen: set[type] = set()
        for cls in classes:
            if cls in seen:
                continue
            seen.add(cls)
            meta = self._build_metadata(cls)
            key = self._key(meta)
            self._classes[key] = cls
            self._metadata[key] = meta
            discovered.append(meta)
        return discovered


    @staticmethod
    def _key(meta: ToolMetadata) -> _RegistryKey:
        return (meta.package, meta.version, meta.module, meta.class_name)

    def _build_metadata(self, cls: type) -> ToolMetadata:
        canonical = getattr(cls, "_bif_canonical_module", cls.__module__)
        package = getattr(cls, "_bif_package", "")
        version = getattr(cls, "_bif_package_version", "")
        try:
            inputs_schema = serialize_input_schema(cls)
        except Exception:
            inputs_schema = {}
        try:
            outputs_schema = serialize_output_schema(cls)
        except Exception:
            outputs_schema = {}
        display_name = (
            getattr(cls, "display_name", None) or cls.__name__
        )
        tags = tuple(getattr(cls, "tags", ()) or ())
        return ToolMetadata(
            package=package,
            version=version,
            module=canonical,
            class_name=cls.__name__,
            inputs_schema=inputs_schema,
            outputs_schema=outputs_schema,
            display_name=display_name,
            tags=tags,
        )

    # -- lookups ------------------------------------------------------------



    def get_class(
        self,
        class_name: str,
        *,
        package: str | None = None,
        version: str | None = None,
        module: str | None = None,
    ) -> type | None:
        """Return a registered class, or ``None`` if not registered.

        When several package versions expose the same class name, callers can
        pass ``package`` / ``version`` / ``module`` to select the intended
        class. A name-only lookup keeps the historical behavior and returns the
        most recently registered matching class.
        """
        match = self._find_key(class_name, package=package, version=version, module=module)
        if match is None:
            return None
        return self._classes[match]




    def get_metadata(
        self,
        class_name: str,
        *,
        package: str | None = None,
        version: str | None = None,
        module: str | None = None,
    ) -> ToolMetadata | None:
        """Return registered :class:`ToolMetadata`, or ``None``."""
        match = self._find_key(class_name, package=package, version=version, module=module)
        if match is None:
            return None
        return self._metadata[match]




    def list_tools(self) -> list[ToolMetadata]:
        """Return all registered tool metadata, in insertion order."""
        return list(self._metadata.values())




    def forget(
        self,
        class_name: str,
        *,
        package: str | None = None,
        version: str | None = None,
        module: str | None = None,
    ) -> None:
        """Drop matching classes from the registry. No-op if none match."""
        for key in list(self._matching_keys(
            class_name,
            package=package,
            version=version,
            module=module,
        )):
            self._classes.pop(key, None)
            self._metadata.pop(key, None)


    def _find_key(
        self,
        class_name: str,
        *,
        package: str | None = None,
        version: str | None = None,
        module: str | None = None,
    ) -> _RegistryKey | None:
        matches = list(self._matching_keys(
            class_name,
            package=package,
            version=version,
            module=module,
        ))
        if not matches:
            return None
        return matches[-1]

    def _matching_keys(
        self,
        class_name: str,
        *,
        package: str | None = None,
        version: str | None = None,
        module: str | None = None,
    ) -> list[_RegistryKey]:
        return [
            key
            for key in self._metadata
            if key[3] == class_name
            and (package is None or key[0] == package)
            and (version is None or key[1] == version)
            and (module is None or key[2] == module)
        ]