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

1#! /usr/bin/env python3 

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

3import sys 

4import subprocess 

5from abc import ABC, abstractmethod 

6from typing import Any 

7 

8from loguru import logger 

9 

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