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

1#!/usr/bin/env python

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

3

4 import sys

5 import json

6 import traceback

7 from io import StringIO

8

9 from loguru import logger

10

11 from .mod_obj import ObjectImporter

12 from .mod_dict import DictModuleImporter

13

14 INIT_IMPORTS = """

15 import os

16 import re

17 import sys

18 import json

19 import time

20 import random

21 import traceback

22"""

23

24 def 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

31 def 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

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