Spaces:
Running
Running
File size: 2,028 Bytes
e70c280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | from __future__ import annotations
import importlib.util
import os
import sys
from pathlib import Path
def _load_module(plugin_dir: Path, relative_path: str, module_name: str):
path = plugin_dir / relative_path
spec = importlib.util.spec_from_file_location(module_name, path)
if spec is None or spec.loader is None:
raise ImportError(f"Cannot load {path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def discover_plugin_celery(plugins_dir: Path) -> tuple[tuple, dict]:
if not plugins_dir.exists():
return (), {}
queues: list = []
routes: dict = {}
for child in sorted(plugins_dir.iterdir()):
if not child.is_dir():
continue
celery_module_path = child / "workers" / "celery.py"
tasks_module_path = child / "workers" / "tasks.py"
if not celery_module_path.exists() or not tasks_module_path.exists():
continue
plugin_dir_str = str(child)
added_plugin_path = plugin_dir_str not in sys.path
if added_plugin_path:
sys.path.insert(0, plugin_dir_str)
try:
safe_name = child.name.replace("-", "_")
celery_mod = _load_module(
child,
"workers/celery.py",
f"brainapi_plugin_{safe_name}_workers_celery",
)
_load_module(
child,
"workers/tasks.py",
f"brainapi_plugin_{safe_name}_workers_tasks",
)
queues.extend(getattr(celery_mod, "QUEUES", ()))
routes.update(getattr(celery_mod, "ROUTES", {}))
finally:
if added_plugin_path and plugin_dir_str in sys.path:
sys.path.remove(plugin_dir_str)
sys.path_importer_cache.pop(plugin_dir_str, None)
return tuple(queues), routes
def default_plugins_dir(project_root: Path) -> Path:
return Path(os.getenv("PLUGINS_DIR", str(project_root / "plugins")))
|