Coverage for aipyapp/interface.py: 53%

1#!/usr/bin/env python

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

3

4 import threading

5 from abc import ABC, abstractmethod

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

7

8 from loguru import logger

9

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

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

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

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

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

77 EventHandler = Callable[[Event], None]

78

79 class EventListener(Protocol):

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

81 ...

82

83 class 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']