import sys import uuid import pkgutil import importlib from hashlib import md5 from pathlib import Path from types import ModuleType from collections import Counter from contextvars import ContextVar from importlib.abc import MetaPathFinder from typing import Set, List, Optional, Iterable from importlib.machinery import PathFinder, SourceFileLoader from .export import _export, Export _current_plugin: ContextVar[Optional[ModuleType]] = ContextVar( "_current_plugin", default=None) _internal_space = ModuleType(__name__ + "._internal") _internal_space.__path__ = [] # type: ignore sys.modules[_internal_space.__name__] = _internal_space _manager_stack: List["PluginManager"] = [] class _NamespaceModule(ModuleType): """Simple namespace module to store plugins.""" @property def __path__(self): return [] def __getattr__(self, name: str): try: return super().__getattr__(name) # type: ignore except AttributeError: if name.startswith("__"): raise raise RuntimeError("Plugin manager not activated!") class _InternalModule(ModuleType): """Internal module for each plugin manager.""" def __init__(self, prefix: str, plugin_manager: "PluginManager"): super().__init__(f"{prefix}.{plugin_manager.internal_id}") self.__plugin_manager__ = plugin_manager @property def __path__(self) -> List[str]: return list(self.__plugin_manager__.search_path) class PluginManager: def __init__(self, namespace: str, plugins: Optional[Iterable[str]] = None, search_path: Optional[Iterable[str]] = None, *, id: Optional[str] = None): self.namespace: str = namespace self.namespace_module: ModuleType = self._setup_namespace(namespace) self.id: str = id or str(uuid.uuid4()) self.internal_id: str = md5( ((self.namespace or "") + self.id).encode()).hexdigest() self.internal_module = self._setup_internal_module(self.internal_id) # simple plugin not in search path self.plugins: Set[str] = set(plugins or []) self.search_path: Set[str] = set(search_path or []) # ensure can be loaded self.list_plugins() def _setup_namespace(self, namespace: str) -> ModuleType: try: module = importlib.import_module(namespace) except ImportError: module = _NamespaceModule(namespace) if "." in namespace: parent = importlib.import_module(namespace.rsplit(".", 1)[0]) setattr(parent, namespace.rsplit(".", 1)[1], module) sys.modules[namespace] = module return module def _setup_internal_module(self, internal_id: str) -> ModuleType: if hasattr(_internal_space, internal_id): raise RuntimeError("Plugin manager already exists!") index = 2 prefix: str = _internal_space.__name__ while True: try: frame = sys._getframe(index) except ValueError: break # check if is called in plugin if "__plugin_name__" not in frame.f_globals: index += 1 continue prefix = frame.f_globals.get("__name__", _internal_space.__name__) break if not prefix.startswith(_internal_space.__name__): prefix = _internal_space.__name__ module = _InternalModule(prefix, self) sys.modules[module.__name__] = module # type: ignore setattr(_internal_space, internal_id, module) return module def __enter__(self): if self in _manager_stack: raise RuntimeError("Plugin manager already activated!") _manager_stack.append(self) return self def __exit__(self, exc_type, exc_value, traceback): try: _manager_stack.pop() except IndexError: pass def search_plugins(self) -> List[str]: return [ module_info.name for module_info in pkgutil.iter_modules(self.search_path) ] def list_plugins(self) -> Set[str]: _pre_managers: List[PluginManager] if self in _manager_stack: _pre_managers = _manager_stack[:_manager_stack.index(self)] else: _pre_managers = _manager_stack[:] _search_path: Set[str] = set() for manager in _pre_managers: _search_path |= manager.search_path if _search_path & self.search_path: raise RuntimeError("Duplicate plugin search path!") _search_plugins = self.search_plugins() c = Counter([*_search_plugins, *self.plugins]) conflict = [name for name, num in c.items() if num > 1] if conflict: raise RuntimeError( f"More than one plugin named {' / '.join(conflict)}!") return set(_search_plugins) | self.plugins def load_plugin(self, name) -> ModuleType: if name in self.plugins: with self: return importlib.import_module(name) if "." in name: raise ValueError("Plugin name cannot contain '.'") with self: return importlib.import_module(f"{self.namespace}.{name}") def load_all_plugins(self) -> List[ModuleType]: return [self.load_plugin(name) for name in self.list_plugins()] def _rewrite_module_name(self, module_name: str) -> Optional[str]: prefix = f"{self.internal_module.__name__}." raw_name = module_name[len(self.namespace) + 1:] if module_name.startswith(self.namespace + ".") else None # dir plugins if raw_name and raw_name.split(".")[0] in self.search_plugins(): return f"{prefix}{raw_name}" # third party plugin or renamed dir plugins elif module_name in self.plugins or module_name.startswith(prefix): return module_name # dir plugins elif module_name in self.search_plugins(): return f"{prefix}{module_name}" return None def _check_absolute_import(self, origin_path: str) -> Optional[str]: if not self.search_path: return paths = set([ *self.search_path, *(str(Path(path).resolve()) for path in self.search_path) ]) for path in paths: try: rel_path = Path(origin_path).relative_to(path) if rel_path.stem == "__init__": return f"{self.internal_module.__name__}." + ".".join( rel_path.parts[:-1]) return f"{self.internal_module.__name__}." + ".".join( rel_path.parts[:-1] + (rel_path.stem,)) except ValueError: continue class PluginFinder(MetaPathFinder): def find_spec(self, fullname: str, path, target): if _manager_stack: index = -1 origin_spec = PathFinder.find_spec(fullname, path, target) while -index <= len(_manager_stack): manager = _manager_stack[index] rel_name = None if origin_spec and origin_spec.origin: rel_name = manager._check_absolute_import( origin_spec.origin) newname = manager._rewrite_module_name(rel_name or fullname) if newname: spec = PathFinder.find_spec( newname, path or [*manager.search_path, *sys.path], target) if spec: spec.loader = PluginLoader( # type: ignore manager, newname, spec.origin) return spec index -= 1 return None class PluginLoader(SourceFileLoader): def __init__(self, manager: PluginManager, fullname: str, path) -> None: self.manager = manager self.loaded = False super().__init__(fullname, path) def create_module(self, spec) -> Optional[ModuleType]: if self.name in sys.modules: self.loaded = True return sys.modules[self.name] # return None to use default module creation return super().create_module(spec) def exec_module(self, module: ModuleType) -> None: if self.loaded: return export = Export() _export_token = _export.set(export) prefix = self.manager.internal_module.__name__ is_dir_plugin = self.name.startswith(prefix + ".") module_name = self.name[len(prefix) + 1:] if is_dir_plugin else self.name _plugin_token = _current_plugin.set(module) setattr(module, "__export__", export) setattr(module, "__plugin_name__", module_name.split(".")[0] if is_dir_plugin else module_name) setattr(module, "__module_name__", module_name) setattr(module, "__module_prefix__", prefix if is_dir_plugin else "") # try: # super().exec_module(module) # except Exception as e: # raise ImportError( # f"Error when executing module {module_name} from {module.__file__}." # ) from e super().exec_module(module) _current_plugin.reset(_plugin_token) _export.reset(_export_token) return sys.meta_path.insert(0, PluginFinder())