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

1#!/usr/bin/env python 

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

3 

4import os 

5import importlib.util 

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

7from pathlib import Path 

8 

9from loguru import logger 

10 

11from .. import __pkgpath__, Plugin, PluginType 

12 

13class 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