From 2909e6e7b9d182e6e9520982923b2b578759c760 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 30 Dec 2024 13:25:25 -0500 Subject: [PATCH] make plugin setup concise --- src/axolotl/integrations/base.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/axolotl/integrations/base.py b/src/axolotl/integrations/base.py index 26f2f8a6f..12a074127 100644 --- a/src/axolotl/integrations/base.py +++ b/src/axolotl/integrations/base.py @@ -223,7 +223,17 @@ def load_plugin(plugin_name: str) -> BasePlugin: module_name, class_name = plugin_name.rsplit(".", 1) # import the module - module = importlib.import_module(module_name) + try: + module = importlib.import_module(module_name) + except ModuleNotFoundError as orig_exc: + try: + if not module_name.startswith("axolotl.integrations."): + module = importlib.import_module("axolotl.integrations." + module_name) + else: + raise orig_exc + except ModuleNotFoundError as exc: + raise orig_exc from exc + # instantiate the class plugin_class = getattr(module, class_name) # create an instance of the class