Coverage for selenible/jupyter.py: 25%

1import base64 

2import inspect 

3import io 

4import shlex 

5from logging import DEBUG, StreamHandler, getLogger 

6from typing import ClassVar 

7 

8import yaml 

9from ipykernel.kernelbase import Kernel 

10from PIL import Image 

11from selenium.webdriver.remote.webelement import WebElement 

12 

13from . import cli 

14from .version import VERSION 

15 

16 

17class SelenibleKernel(Kernel): 

18 log = getLogger("selenible") 

19 implementation = "Selenible" 

20 implementation_version = "0.0.1" 

21 language = "selenible" 

22 language_version = VERSION 

23 language_info: ClassVar = { 

24 "name": "Selenible", 

25 "mimetype": "text/yaml", 

26 "file_extension": ".yaml", 

27 } 

28 banner = "Selenible kernel" 

29 driver_name = "phantom" 

30 extensions: ClassVar = [] 

31 

32 def __init__(self, **kwargs): 

33 super().__init__(**kwargs) 

34 self._drv = None 

35 self.thumbnail = None 

36 self.log.setLevel(DEBUG) 

37 self.log.info("kernel started") 

38 

39 @property 

40 def drv(self): 

41 if self._drv is None: 

42 drvcls = cli.loadmodules(self.driver_name, self.extensions) 

43 self.log.info( 

44 "driver: cls=%s, name=%s, exts=%s", 

45 drvcls, 

46 self.driver_name, 

47 self.extensions, 

48 ) 

49 self._drv = drvcls() 

50 drvlog = self._drv.log 

51 self.logio = io.StringIO() 

52 drvlog.addHandler(StreamHandler(self.logio)) 

53 return self._drv 

54 

55 def do_shutdown(self, restart): 

56 self.log.info("kernel finished") 

57 del self._drv 

58 self._drv = None 

59 

60 def cmd_driver(self, args): 

61 "set driver: phantom, chrome, firefox, etc..." 

62 self.log.info("driver: %s -> %s", self.driver_name, args[0]) 

63 self.driver_name = args[0] 

64 

65 def cmd_module(self, args): 

66 "load modules" 

67 self.extensions = args 

68 

69 def cmd_shutdown(self, args): 

70 "shutdown driver" 

71 del self._drv 

72 self._drv = None 

73 

74 def cmd_loglevel(self, args): 

75 "set log level" 

76 self.drv.log.setLevel(args[0]) 

77 

78 def cmd_thumbnail(self, args): 

79 "set thumbnail size" 

80 self.thumbnail = tuple(int(f) for f in args[:2]) 

81 

82 def cmd_help(self, args): 

83 "show this help" 

84 cmds = filter(lambda f: f.startswith("cmd_"), dir(self)) 

85 cmds = sorted(filter(lambda f: callable(getattr(self, f)), cmds)) 

86 cmds = [(x.split("_", 1)[1], inspect.getdoc(getattr(self, x))) for x in cmds] 

87 txt = "\n".join(["{}: {}".format(*x) for x in cmds]) 

88 self.send_response(self.iopub_socket, "stream", {"name": "stdout", "text": txt}) 

89 

90 def do_execute( 

91 self, code, silent, store_history=True, user_expressions=None, allow_stdin=False 

92 ): 

93 if code.startswith("%"): 

94 token = shlex.split(code) 

95 cmd = token[0].lstrip("%") 

96 args = token[1:] 

97 if hasattr(self, "cmd_" + cmd): 

98 fn = getattr(self, "cmd_" + cmd) 

99 if callable(fn): 

100 fn(args) 

101 else: 

102 self.cmd_help() 

103 else: 

104 return { 

105 "status": "error", 

106 "ename": "NotFound", 

107 "evalue": "not found", 

108 "traceback": [], 

109 } 

110 return {"status": "ok", "execution_count": self.execution_count} 

111 v = yaml.safe_load(code) 

112 self.log.info("yaml: %s", v) 

113 if not silent: 

114 stream_content = { 

115 "name": "stdout", 

116 "text": yaml.dump(v, default_flow_style=False), 

117 } 

118 self.send_response(self.iopub_socket, "stream", stream_content) 

119 if isinstance(v, (list, tuple)): 

120 res = self.drv.run(v) 

121 elif isinstance(v, dict): 

122 res = self.drv.run([v]) 

123 elif isinstance(v, str): 

124 res = self.drv.run([{v: None}]) 

125 else: 

126 raise Exception(f"invalid type: {type(v)} : {v}") 

127 

128 logstr = self.logio.getvalue() 

129 self.logio.seek(0) 

130 self.logio.truncate(0) 

131 if logstr != "": 

132 stream_content = { 

133 "data": {"text/plain": logstr}, 

134 "execution_count": self.execution_count, 

135 } 

136 self.send_response(self.iopub_socket, "execute_result", stream_content) 

137 if not isinstance(res, (str, dict, list, tuple)): 

138 self.log.info("not json serializeable?: %s", res) 

139 ress = str(res) 

140 else: 

141 ress = res 

142 if res is not None: 

143 stream_content = { 

144 "data": {"text/plain": ress}, 

145 "execution_count": self.execution_count, 

146 } 

147 self.send_response(self.iopub_socket, "execute_result", stream_content) 

148 if isinstance(res, WebElement): 

149 imgdata = res.screenshot_as_png 

150 else: 

151 imgdata = self.drv.saveshot() 

152 img = Image.open(io.BytesIO(imgdata)) 

153 if self.thumbnail is not None: 

154 olen = len(imgdata) 

155 img.thumbnail(self.thumbnail, Image.ANTIALIAS) 

156 buf = io.BytesIO() 

157 img.save(buf, format="png") 

158 imgdata = buf.getvalue() 

159 self.log.info("datasize: %d -> %d", olen, len(imgdata)) 

160 imgdict = { 

161 "data": { 

162 "image/png": base64.b64encode(imgdata).decode("ascii"), 

163 }, 

164 "metadata": { 

165 "image/png": { 

166 "width": img.size[0], 

167 "height": img.size[1], 

168 } 

169 }, 

170 } 

171 self.send_response(self.iopub_socket, "display_data", imgdict) 

172 self.log.info("image: %s", img.size) 

173 return {"status": "ok", "execution_count": self.execution_count} 

174 

175 

176def main(): 

177 from ipykernel.kernelapp import IPKernelApp 

178 

179 IPKernelApp.launch_instance(kernel_class=SelenibleKernel) 

180 

181 

182if __name__ == "__main__": 182 ↛ 183line 182 didn't jump to line 183 because the condition on line 182 was never true

183 main()