fix: 让自动获取插件名更稳定 (#216)

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/lang/zh-CN/
 
 ## [Unreleased]
 
+### Fixed
+
+- 让自动获取插件名更稳定
+
 ## [1.2.0] - 2024-02-28
 
 ### Added

nonebot_plugin_datastore/providers/database.py

@@ -8,7 +8,7 @@
 from ..plugin import get_plugin_data
 from . import ConfigProvider, KeyNotFoundError
 
-plugin_data = get_plugin_data()
+plugin_data = get_plugin_data("nonebot_plugin_datastore")
 
 
 class ConfigModel(plugin_data.Model):

nonebot_plugin_datastore/utils.py

@@ -1,50 +1,32 @@
 import importlib
 import inspect
-from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
-import pygtrie
-from nonebot import get_loaded_plugins
-
-if TYPE_CHECKING:
-    from nonebot.plugin import Plugin
+from nonebot import get_plugin_by_module_name
 
 
 def get_caller_plugin_name() -> str:
     """获取当前函数调用者所在的插件名
 
     尝试自动获取调用者所在的插件名
     """
-    name = None
-    if frame := inspect.currentframe():
-        # 因为是在插件内部调用，所以调用栈为
-        # 1. 当前函数
-        # 2. 调用者函数
-        # 3. 调用者所在的插件
-        # 需要往回跳两次
-        frame = frame.f_back.f_back  # type: ignore
-        if not frame:
-            raise ValueError("无法找到调用者")  # pragma: no cover
-
-        module_name = frame.f_locals["__name__"]
-        plugin = _get_plugin_by_module_name(module_name)
-        if plugin:
-            name = plugin.name
-
-    if not name:
-        raise ValueError("自动获取插件名失败")
-
-    return name
-
-
-@lru_cache
-def _get_plugin_by_module_name(module_name: str) -> Optional["Plugin"]:
-    """通过模块名获取插件"""
-    t = pygtrie.StringTrie(separator=".")
-    for plugin in get_loaded_plugins():
-        t[plugin.module_name] = plugin
-    plugin = t.longest_prefix(module_name).value
-    return plugin
+    frame = inspect.currentframe()
+    if frame is None:
+        raise ValueError("无法获取当前栈帧")
+
+    while frame := frame.f_back:
+        module_name = (module := inspect.getmodule(frame)) and module.__name__
+        if not module_name:
+            raise ValueError("无法找到调用者")
+
+        if module_name.split(".", maxsplit=1)[0] == "nonebot_plugin_datastore":
+            continue
+
+        plugin = get_plugin_by_module_name(module_name)
+        if plugin and plugin.id_ != "nonebot_plugin_datastore":
+            return plugin.name
+
+    raise ValueError("自动获取插件名失败")
 
 
 def resolve_dot_notation(

tests/test_plugin.py

@@ -6,12 +6,14 @@ async def test_get_plugin_data_failed(app: App):
     """获取插件数据失败"""
     from nonebot_plugin_datastore import get_plugin_data
 
-    # f_locals 中没有 __name__
-    with pytest.raises(KeyError):
+    # 不在插件中调用
+    # 挺奇妙的，如果用 pytest 跑会报无法找到调用者
+    # 但是 vscode 调试中跑就会报自动获取插件名失败
+    with pytest.raises(ValueError, match=r"无法找到调用者|自动获取插件名失败"):
         get_plugin_data()
 
     # 没有加载插件直接使用
-    with pytest.raises(ValueError, match="自动获取插件名失败"):
+    with pytest.raises(ValueError, match="无法找到调用者"):
         import tests.example.plugin1  # noqa: F401