Coverage for aipyapp/interface.py: 53%

1#!/usr/bin/env python 

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

3 

4import threading 

5from abc import ABC, abstractmethod 

6from typing import Callable, Any, Dict, List, Optional, Protocol 

7 

8from loguru import logger 

9 

10class Trackable(ABC): 

11 """可追踪对象接口""" 

12 

13 @abstractmethod 

14 def get_checkpoint(self) -> Any: 

15 """获取当前检查点状态""" 

16 pass 

17 

18 @abstractmethod 

19 def restore_to_checkpoint(self, checkpoint: Optional[Any]): 

20 """恢复到指定检查点，None表示恢复到初始状态""" 

21 pass 

22 

23class Runtime(ABC): 

24 @abstractmethod 

25 def install_packages(self, packages): 

26 pass 

27 

28 @abstractmethod 

29 def getenv(self, name, desc=None): 

30 pass 

31 

32class ConsoleInterface(ABC): 

33 @abstractmethod 

34 def print(self, *args, sep=' ', end='\n', file=None, flush=False): 

35 pass 

36 

37 @abstractmethod 

38 def input(self, prompt=''): 

39 pass 

40 

41 @abstractmethod 

42 def status(self, msg): 

43 pass 

44 

45class Stoppable(): 

46 def __init__(self): 

47 super().__init__() 

48 self._stop_event = threading.Event() 

49 

50 def on_stop(self): 

51 pass 

52 

53 def stop(self): 

54 self._stop_event.set() 

55 self.on_stop() 

56 

57 def is_stopped(self): 

58 return self._stop_event.is_set() 

59 

60 def wait(self, timeout=None): 

61 return self._stop_event.wait(timeout) 

62 

63 def reset(self): 

64 self._stop_event.clear() 

65 

66class Event: 

67 def __init__(self, name: str, **data): 

68 self.name = name 

69 self.data = data 

70 

71 def __str__(self): 

72 return f"{self.name}: {self.data}" 

73 

74 def __getattr__(self, name: str): 

75 return self.data.get(name) 

76 

77EventHandler = Callable[[Event], None] 

78 

79class EventListener(Protocol): 

80 def get_handlers(self) -> Dict[str, EventHandler]: 

81 ... 

82 

83class EventBus: 

84 def __init__(self): 

85 super().__init__() 

86 self._listeners: Dict[str, List[EventHandler]] = {} 

87 self._eb_logger = logger.bind(src=self.__class__.__name__) 

88 

89 def on_event(self, event_name: str, handler: EventHandler): 

90 self._listeners.setdefault(event_name, []).append(handler) 

91 

92 def add_listener(self, obj: EventListener): 

93 count = 0 

94 for event_name, handler in obj.get_handlers().items(): 

95 self.on_event(event_name, handler) 

96 count += 1 

97 self._eb_logger.info(f"Registered {count} events for {obj.__class__.__name__}") 

98 

99 def emit(self, event_name: str, **kwargs): 

100 event = Event(event_name, **kwargs) 

101 for handler in self._listeners.get(event_name, []): 

102 try: 

103 handler(event) 

104 except Exception as e: 

105 self._eb_logger.exception(f"Error emitting event {event_name}") 

106 return event 

107 

108 

109__all__ = ['Trackable', 'Runtime', 'Stoppable', 'Event', 'EventHandler', 'EventListener', 'EventBus']