Source code for core.registry

"""Probe registry — discovers and runs probes by pattern matching.

Supports two discovery mechanisms:
1. Built-in: auto-discover probes under probes/ package
2. Plugins: external packages register via entry_points group 'metapod.probes'

To create a plugin:
  # In your package's pyproject.toml:
  [project.entry-points."metapod.probes"]
  my_source = "my_package.probes"
"""

from __future__ import annotations

import fnmatch
import importlib
import pkgutil
from pathlib import Path

from core.base_probe import BaseProbe

# Global registry
_PROBES: dict[str, type[BaseProbe]] = {}
_discovered = False




def register(cls: type[BaseProbe]) -> type[BaseProbe]:
    """Decorator to register a probe class."""
    name = cls.probe_name()
    if name in _PROBES:
        import logging
        logging.getLogger("metapod").warning("Duplicate probe registration: %s (overwriting)", name)
    _PROBES[name] = cls
    return cls





def get_probe(name: str) -> type[BaseProbe] | None:
    return _PROBES.get(name)





def list_probes(pattern: str = "*") -> list[str]:
    """List registered probe names matching a glob pattern."""
    return sorted(name for name in _PROBES if fnmatch.fnmatch(name, pattern))





def get_probes(pattern: str = "*") -> list[type[BaseProbe]]:
    """Get probe classes matching a glob pattern."""
    return [_PROBES[name] for name in list_probes(pattern)]





def discover_probes():
    """Auto-discover all probes: built-in + plugins."""
    global _discovered
    if _discovered:
        return
    _discovered = True

    # 1. Built-in probes under probes/ package
    probes_dir = Path(__file__).parent.parent / "probes"
    if probes_dir.exists():
        for source_type in probes_dir.iterdir():
            if not source_type.is_dir() or source_type.name.startswith("_"):
                continue
            for category in source_type.iterdir():
                if not category.is_dir() or category.name.startswith("_"):
                    continue
                for module_info in pkgutil.iter_modules([str(category)]):
                    if module_info.name.startswith("_"):
                        continue
                    module_path = f"probes.{source_type.name}.{category.name}.{module_info.name}"
                    try:
                        importlib.import_module(module_path)
                    except Exception as e:
                        import logging
                        logging.getLogger("metapod").debug("Failed to load %s: %s", module_path, e)

    # 2. External plugins via entry_points
    try:
        if hasattr(importlib.metadata, "entry_points"):
            eps = importlib.metadata.entry_points()
            # Python 3.12+ returns a dict, 3.10-3.11 returns SelectableGroups
            if hasattr(eps, "select"):
                plugin_eps = eps.select(group="metapod.probes")
            elif isinstance(eps, dict):
                plugin_eps = eps.get("metapod.probes", [])
            else:
                plugin_eps = getattr(eps, "metapod.probes", [])

            for ep in plugin_eps:
                try:
                    ep.load()
                except Exception as e:
                    import logging
                    logging.getLogger("metapod").debug("Plugin load failed %s: %s", ep.name, e)
    except Exception:
        pass