Coverage for aipyapp/aipy/plugins.py: 21%

1#!/usr/bin/env python

2# -*- coding: utf-8 -*-

3

4 import os

5 import importlib.util

6 from typing import Dict, Any, List, Union, Optional, Type

7 from pathlib import Path

8

9 from loguru import logger

10

11 from .. import __pkgpath__, Plugin, PluginType

12

13 class PluginManager:

14 """Plugin manager class."""

15 FILE_PATTERN = "p_*.py"

16

17 def __init__(self):

18 self.sys_plugin_dir = os.path.join(__pkgpath__, 'plugins')

19 self.plugin_directories: List[Path] = []

20 self._plugins: Dict[str, Any] = {}

21 self.logger = logger.bind(src=self.__class__.__name__)

22 self.add_plugin_directory(self.sys_plugin_dir)

23

24 def __iter__(self):

25 return iter(self._plugins.values())

26

27 def __getitem__(self, name: str) -> Plugin:

28 return self._plugins[name]

29

30 def __len__(self) -> int:

31 return len(self._plugins)

32

33 def add_plugin_directory(self, directory: Union[str, Path]):

34 """Add a plugin directory to the plugin manager."""

35 dir_path = Path(directory)

36 if not dir_path.is_dir():

37 self.logger.warning(f"Plugin directory {directory} does not exist")

38 return False

39 if dir_path in self.plugin_directories:

40 self.logger.warning(f"Plugin directory {directory} already added")

41 return False

42 self.plugin_directories.append(dir_path)

43 self.logger.info(f"Added plugin directory {directory}")

44 return True

45

46 def _discover_plugins(self) -> List[Path]:

47 """Discover plugins in the plugin directories."""

48 plugin_files = []

49 for plugin_dir in self.plugin_directories:

50 if not os.path.exists(plugin_dir):

51 self.logger.warning(f"Plugin directory {plugin_dir} does not exist")

52 continue

53 for fname in plugin_dir.glob(self.FILE_PATTERN):

54 plugin_files.append(fname)

55 return plugin_files

56

57 def _load_plugins(self, filepath: Path) -> List[Plugin]:

58 """Load plugins from a file."""

59 import sys

60

61 name = filepath.stem

62 plugin_dir = filepath.parent

63

64 # 临时添加插件目录到 sys.path

65 original_path = sys.path.copy()

66 if str(plugin_dir) not in sys.path:

67 sys.path.insert(0, str(plugin_dir))

68

69 try:

70 spec = importlib.util.spec_from_file_location(name, filepath)

71 module = importlib.util.module_from_spec(spec)

72 spec.loader.exec_module(module)

73

74 plugins = []

75 for cls in module.__dict__.values():

76 if isinstance(cls, type) and issubclass(cls, Plugin) and cls.name:

77 plugins.append(cls)

78 return plugins

79 finally:

80 # 恢复原始 sys.path

81 sys.path = original_path

82

83 def _register_plugin(self, plugin: Plugin) -> bool:

84 """Register a plugin."""

85 plugin_name = plugin.name

86 if not plugin_name:

87 self.logger.warning(f"Plugin {plugin.__class__.__name__} has no name")

88 return False

89 if plugin_name in self._plugins:

90 self.logger.warning(f"Plugin {plugin_name} already registered")

91 return False

92 self._plugins[plugin_name] = plugin

93 return True

94

95 def load_all_plugins(self):

96 """Load all plugins."""

97 success_count = 0

98 plugin_files = self._discover_plugins()

99 for plugin_file in plugin_files:

100 plugins = self._load_plugins(plugin_file)

101 for plugin in plugins:

102 if self._register_plugin(plugin):

103 success_count += 1

104 self.logger.info(f"Loaded plugin {plugin.name} from {plugin_file}")

105 self.logger.info(f"Loaded {success_count} plugins")

106

107 def create_task_plugin(self, name: str, plugin_config: Dict[str, Any] = None) -> Optional[Plugin]:

108 """Create a plugin by name."""

109 plugin_cls = self._plugins.get(name)

110 if not plugin_cls:

111 self.logger.warning(f"Plugin {name} not found")

112 return

113

114 if plugin_cls.get_type() != PluginType.TASK:

115 self.logger.warning(f"Plugin {name} is not a task plugin")

116 return None

117

118 plugin = plugin_cls(plugin_config)

119 try:

120 plugin.init()

121 except Exception as e:

122 self.logger.error(f"Failed to initialize plugin {name}: {e}")

123 return None

124 return plugin

125

126 def get_task_plugins(self) -> List[Type[Plugin]]:

127 """Get all task plugins."""

128 return [plugin for plugin in self._plugins.values() if plugin.get_type() == PluginType.TASK]

129

130 def get_display_plugins(self) -> List[Type[Plugin]]:

131 """Get all display plugins."""

132 return [plugin for plugin in self._plugins.values() if plugin.get_type() == PluginType.DISPLAY]

133