Coverage for aipyapp/exec/executor.py: 29%

1#!/usr/bin/env python

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

3

4 import traceback

5 from typing import Optional

6

7 from loguru import logger

8

9 from ..interface import Trackable

10 from .python import PythonRuntime, PythonExecutor

11 from .html import HtmlExecutor

12 from .prun import BashExecutor, PowerShellExecutor, AppleScriptExecutor, NodeExecutor

13

14 EXECUTORS = {executor.name: executor for executor in [

15 PythonExecutor,

16 HtmlExecutor,

17 BashExecutor,

18 PowerShellExecutor,

19 AppleScriptExecutor,

20 NodeExecutor

21]}

22

23 class BlockExecutor(Trackable):

24 def __init__(self):

25 self.history = []

26 self.executors = {}

27 self.runtimes = {}

28 self.log = logger.bind(src='block_executor')

29

30 def _set_runtime(self, lang, runtime):

31 if lang not in self.runtimes:

32 if lang not in EXECUTORS:

33 self.log.warning(f'No executor found for {lang}')

34 self.runtimes[lang] = runtime

35 self.log.info(f'Registered runtime for {lang}: {runtime}')

36 else:

37 self.log.warning(f'Runtime for {lang} already registered: {self.runtimes[lang]}')

38

39 def set_python_runtime(self, runtime):

40 assert isinstance(runtime, PythonRuntime), "Expected a PythonRuntime instance"

41 self._set_runtime('python', runtime)

42

43 def get_executor(self, block):

44 lang = block.get_lang()

45 if lang in self.executors:

46 return self.executors[lang]

47

48 if lang not in EXECUTORS:

49 self.log.warning(f'No executor found for {lang}')

50 return None

51

52 runtime = self.runtimes.get(lang)

53 executor = EXECUTORS[lang](runtime)

54 self.executors[lang] = executor

55 self.log.info(f'Registered executor for {lang}: {executor}')

56 return executor

57

58 def __call__(self, block):

59 self.log.info(f'Exec: {block}')

60 history = {}

61 executor = self.get_executor(block)

62 if executor:

63 try:

64 result = executor(block)

65 except Exception as e:

66 result = {'errstr': str(e), 'traceback': traceback.format_exc()}

67 else:

68 result = {'stderr': f'Exec: Ignore unsupported block: {block}'}

69

70 history['block'] = block

71 history['result'] = result

72 self.history.append(history)

73 return result

74

75 def get_state(self):

76 """获取需要持久化的状态数据"""

77 return self.history.copy()

78

79 def restore_state(self, runner_data):

80 """从运行历史数据恢复状态"""

81 self.history.clear()

82 if runner_data:

83 self.history = runner_data.copy()

84

85

86 def clear(self):

87 self.history.clear()

88

89 # Trackable接口实现

90 def get_checkpoint(self) -> int:

91 """获取当前检查点状态 - 返回执行历史长度"""

92 return len(self.history)

93

94 def restore_to_checkpoint(self, checkpoint: Optional[int]):

95 """恢复到指定检查点"""

96 if checkpoint is None:

97 # 恢复到初始状态

98 self.clear()

99 else:

100 # 恢复到指定长度

101 if checkpoint < len(self.history):

102 self.history = self.history[:checkpoint]