Coverage for aipyapp/exec/python/runtime.py: 35%

1#! /usr/bin/env python3

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

3 import sys

4 import subprocess

5 from abc import ABC, abstractmethod

6 from typing import Any

7

8 from loguru import logger

9

10 class PythonRuntime(ABC):

11 def __init__(self, envs=None):

12 self.envs = envs or {}

13 self.packages = set()

14 self.session = {}

15 self.block_states = {}

16 self.current_state = {}

17 self.block = None

18 self.log = logger.bind(src='runtime')

19

20 def start_block(self, block):

21 """开始一个新的代码块执行"""

22 self.current_state = {}

23 self.block_states[block.name] = self.current_state

24 self.block = block

25

26 def set_state(self, success: bool, **kwargs) -> None:

27 """

28 Set the state of the current code block

29

30 Args:

31 success: Whether the code block is successful

32 **kwargs: Other state values

33

34 Example:

35 set_state(success=True, result="Hello, world!")

36 set_state(success=False, error="Error message")

37 """

38 self.current_state['success'] = success

39 self.current_state.update(kwargs)

40

41 def get_block_state(self, block_name: str) -> Any:

42 """

43 Get the state of code block by name

44

45 Args:

46 block_name: The name of the code block

47

48 Returns:

49 Any: The state of the code block

50

51 Example:

52 state = get_block_state("code_block_name")

53 if state.get("success"):

54 print(state.get("result"))

55 else:

56 print(state.get("error"))

57 """

58 return self.block_states.get(block_name)

59

60 def set_persistent_state(self, **kwargs) -> None:

61 """

62 Set the state of the current code block in the session

63

64 Args:

65 **kwargs: The state values

66 """

67 self.session.update(kwargs)

68 self.block.add_dep('set_state', list(kwargs.keys()))

69

70 def get_persistent_state(self, key: str) -> Any:

71 """

72 Get the state of the current code block in the session

73

74 Args:

75 key: The key of the state

76

77 Returns:

78 Any: The state of the code block

79 """

80 self.block.add_dep('get_state', key)

81 return self.session.get(key)

82

83 def set_env(self, name, value, desc):

84 self.envs[name] = (value, desc)

85

86 def ensure_packages(self, *packages, upgrade=False, quiet=False):

87 if not packages:

88 return True

89

90 packages = list(set(packages) - self.packages)

91 if not packages:

92 return True

93

94 cmd = [sys.executable, "-m", "pip", "install"]

95 if upgrade:

96 cmd.append("--upgrade")

97 if quiet:

98 cmd.append("-q")

99 cmd.extend(packages)

100

101 try:

102 subprocess.check_call(cmd)

103 self.packages.update(packages)

104 return True

105 except subprocess.CalledProcessError:

106 self.log.error("依赖安装失败: {}", " ".join(packages))

107

108 return False

109

110 def ensure_requirements(self, path="requirements.txt", **kwargs):

111 with open(path) as f:

112 reqs = [line.strip() for line in f if line.strip() and not line.startswith("#")]

113 return self.ensure_packages(*reqs, **kwargs)

114

115 @abstractmethod

116 def install_packages(self, *packages: str):

117 pass

118

119 @abstractmethod

120 def get_env(self, name: str, default: Any = None, *, desc: str = None) -> Any:

121 pass

122

123 @abstractmethod

124 def show_image(self, path: str = None, url: str = None) -> None:

125 pass

126

127 @abstractmethod

128 def input(self, prompt: str = '') -> str:

129 pass