Coverage for aipyapp/exec/python/executor.py: 21%

1#!/usr/bin/env python 

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

3 

4import sys 

5import json 

6import traceback 

7from io import StringIO 

8 

9from loguru import logger 

10 

11from .mod_obj import ObjectImporter 

12from .mod_dict import DictModuleImporter 

13 

14INIT_IMPORTS = """ 

15import os 

16import re 

17import sys 

18import json 

19import time 

20import random 

21import traceback 

22""" 

23 

24def is_json_serializable(obj): 

25 try: 

26 json.dumps(obj, ensure_ascii=False, default=str) 

27 return True 

28 except (TypeError, OverflowError): 

29 return False 

30 

31def diff_dicts(dict1, dict2): 

32 diff = {} 

33 for key, value in dict1.items(): 

34 if key not in dict2: 

35 diff[key] = value 

36 continue 

37 

38 try: 

39 if value != dict2[key]: 

40 diff[key] = value 

41 except Exception: 

42 pass 

43 return diff 

44 

45class PythonExecutor(): 

46 name = 'python' 

47 

48 def __init__(self, runtime): 

49 self.runtime = runtime 

50 self.log = logger.bind(src='PythonExecutor') 

51 self._globals = {'__name__': '__main__', 'input': self.runtime.input} 

52 self.block_importer = DictModuleImporter() 

53 self.runtime_importer = ObjectImporter({'utils': runtime}) 

54 exec(INIT_IMPORTS, self._globals) 

55 

56 def __repr__(self): 

57 return "<PythonExecutor>" 

58 

59 @property 

60 def globals(self): 

61 return self._globals 

62 

63 def __call__(self, block): 

64 result = {} 

65 try: 

66 co = compile(block.code, block.abs_path or block.name, 'exec') 

67 except SyntaxError as e: 

68 result['errstr'] = f"Syntax error: {str(e)}" 

69 result['traceback'] = traceback.format_exc() 

70 return result 

71 

72 runtime = self.runtime 

73 old_stdout, old_stderr = sys.stdout, sys.stderr 

74 captured_stdout = StringIO() 

75 captured_stderr = StringIO() 

76 sys.stdout, sys.stderr = captured_stdout, captured_stderr 

77 gs = self._globals.copy() 

78 runtime.start_block(block) 

79 try: 

80 with self.block_importer, self.runtime_importer: 

81 exec(co, gs) 

82 self.block_importer.add_module(block.name, co) 

83 except (SystemExit, Exception) as e: 

84 self.runtime.set_state(success=False, error=str(e)) 

85 self.log.error(f"Error in code block {block.name}: {str(e)}") 

86 result['errstr'] = str(e) 

87 result['traceback'] = traceback.format_exc() 

88 finally: 

89 sys.stdout = old_stdout 

90 sys.stderr = old_stderr 

91 

92 s = captured_stdout.getvalue().strip() 

93 if s: result['stdout'] = s if is_json_serializable(s) else '<filtered: cannot json-serialize>' 

94 s = captured_stderr.getvalue().strip() 

95 if s: result['stderr'] = s if is_json_serializable(s) else '<filtered: cannot json-serialize>' 

96 

97 vars = runtime.current_state 

98 if vars: 

99 result['__state__'] = self.filter_result(vars) 

100 

101 return result 

102 

103 def filter_result(self, vars): 

104 if isinstance(vars, dict): 

105 ret = {} 

106 for key in vars.keys(): 

107 if key in self.runtime.envs: 

108 ret[key] = '<masked>' 

109 else: 

110 ret[key] = self.filter_result(vars[key]) 

111 elif isinstance(vars, list): 

112 ret = [self.filter_result(v) for v in vars] 

113 else: 

114 ret = vars if is_json_serializable(vars) else '<filtered>' 

115 return ret