Coverage for /builds/debichem-team/python-ase/ase/visualize/viewers.py: 82.93%

1""" 

2Module for managing viewers 

3 

4View plugins can be registered through the entrypoint system in with the 

5following in a module, such as a `viewer.py` file: 

6 

7```python3 

8VIEWER_ENTRYPOINT = ExternalViewer( 

9 desc="Visualization using <my package>", 

10 module="my_package.viewer" 

11) 

12``` 

13 

14Where module `my_package.viewer` contains a `view_my_viewer` function taking 

15and `ase.Atoms` object as the first argument, and also `**kwargs`. 

16 

17Then ones needs to register an entry point in `pyproject.toml` with 

18 

19```toml 

20[project.entry-points."ase.visualize"] 

21my_viewer = "my_package.viewer:VIEWER_ENTRYPOINT" 

22``` 

23 

24After this, call to `ase.visualize.view(atoms, viewer='my_viewer')` will be 

25forwarded to `my_package.viewer.view_my_viewer` function. 

26 

27""" 

28import pickle 

29import subprocess 

30import sys 

31import tempfile 

32import warnings 

33from contextlib import contextmanager 

34from importlib import import_module 

35from importlib.metadata import entry_points 

36from io import BytesIO 

37from pathlib import Path 

38 

39from ase.io import write 

40from ase.io.formats import ioformats 

41from ase.utils.plugins import ExternalViewer 

42 

43 

44class UnknownViewerError(Exception): 

45 """The view tyep is unknown""" 

46 

47 

48class AbstractViewer: 

49 def view(self, *args, **kwargss): 

50 raise NotImplementedError 

51 

52 

53class PyViewer(AbstractViewer): 

54 def __init__(self, name: str, desc: str, module_name: str): 

55 """ 

56 Instantiate an viewer 

57 """ 

58 self.name = name 

59 self.desc = desc 

60 self.module_name = module_name 

61 

62 def _viewfunc(self): 

63 """Return the function used for viewing the atoms""" 

64 return getattr(self.module, "view_" + self.name, None) 

65 

66 @property 

67 def module(self): 

68 try: 

69 return import_module(self.module_name) 

70 except ImportError as err: 

71 raise UnknownViewerError( 

72 f"Viewer not recognized: {self.name}. Error: {err}" 

73 ) from err 

74 

75 def view(self, atoms, *args, **kwargs): 

76 return self._viewfunc()(atoms, *args, **kwargs) 

77 

78 

79class CLIViewer(AbstractViewer): 

80 """Generic viewer for""" 

81 

82 def __init__(self, name, fmt, argv): 

83 self.name = name 

84 self.fmt = fmt 

85 self.argv = argv 

86 

87 @property 

88 def ioformat(self): 

89 return ioformats[self.fmt] 

90 

91 @contextmanager 

92 def mktemp(self, atoms, data=None): 

93 ioformat = self.ioformat 

94 suffix = "." + ioformat.extensions[0] 

95 

96 if ioformat.isbinary: 

97 mode = "wb" 

98 else: 

99 mode = "w" 

100 

101 with tempfile.TemporaryDirectory(prefix="ase-view-") as dirname: 

102 # We use a tempdir rather than a tempfile because it's 

103 # less hassle to handle the cleanup on Windows (files 

104 # cannot be open on multiple processes). 

105 path = Path(dirname) / f"atoms{suffix}" 

106 with path.open(mode) as fd: 

107 if data is None: 

108 write(fd, atoms, format=self.fmt) 

109 else: 

110 write(fd, atoms, format=self.fmt, data=data) 

111 yield path 

112 

113 def view_blocking(self, atoms, data=None): 

114 with self.mktemp(atoms, data) as path: 

115 subprocess.check_call(self.argv + [str(path)]) 

116 

117 def view( 

118 self, 

119 atoms, 

120 data=None, 

121 repeat=None, 

122 **kwargs, 

123 ): 

124 """Spawn a new process in which to open the viewer.""" 

125 if repeat is not None: 

126 atoms = atoms.repeat(repeat) 

127 

128 proc = subprocess.Popen( 

129 [sys.executable, "-m", "ase.visualize.viewers"], 

130 stdin=subprocess.PIPE 

131 ) 

132 

133 pickle.dump((self, atoms, data), proc.stdin) 

134 proc.stdin.close() 

135 return proc 

136 

137 

138VIEWERS = {} 

139 

140 

141def _pipe_to_ase_gui(atoms, repeat, **kwargs): 

142 buf = BytesIO() 

143 write(buf, atoms, format="traj") 

144 

145 args = [sys.executable, "-m", "ase", "gui", "-"] 

146 if repeat: 

147 args.append("--repeat={},{},{}".format(*repeat)) 

148 

149 proc = subprocess.Popen(args, stdin=subprocess.PIPE) 

150 proc.stdin.write(buf.getvalue()) 

151 proc.stdin.close() 

152 return proc 

153 

154 

155def define_viewer( 

156 name, desc, *, module=None, cli=False, fmt=None, argv=None, external=False 

157): 

158 if not external: 

159 if module is None: 

160 module = name 

161 module = "ase.visualize." + module 

162 if cli: 

163 fmt = CLIViewer(name, fmt, argv) 

164 else: 

165 if name == "ase": 

166 # Special case if the viewer is named `ase` then we use 

167 # the _pipe_to_ase_gui as the viewer method 

168 fmt = PyViewer(name, desc, module_name=None) 

169 fmt.view = _pipe_to_ase_gui 

170 else: 

171 fmt = PyViewer(name, desc, module_name=module) 

172 VIEWERS[name] = fmt 

173 return fmt 

174 

175 

176def define_external_viewer(entry_point): 

177 """Define external viewer""" 

178 

179 viewer_def = entry_point.load() 

180 if entry_point.name in VIEWERS: 

181 raise ValueError(f"Format {entry_point.name} already defined") 

182 if not isinstance(viewer_def, ExternalViewer): 

183 raise TypeError( 

184 "Wrong type for registering external IO formats " 

185 f"in format {entry_point.name}, expected " 

186 "ExternalViewer" 

187 ) 

188 define_viewer(entry_point.name, **viewer_def._asdict(), 

189 external=True) 

190 

191 

192def register_external_viewer_formats(group): 

193 if hasattr(entry_points(), "select"): 

194 viewer_entry_points = entry_points().select(group=group) 

195 else: 

196 viewer_entry_points = entry_points().get(group, ()) 

197 

198 for entry_point in viewer_entry_points: 

199 try: 

200 define_external_viewer(entry_point) 

201 except Exception as exc: 

202 warnings.warn( 

203 "Failed to register external " 

204 f"Viewer {entry_point.name}: {exc}" 

205 ) 

206 

207 

208define_viewer("ase", "View atoms using ase gui.") 

209define_viewer("ngl", "View atoms using nglview.") 

210define_viewer("mlab", "View atoms using matplotlib.") 

211define_viewer("sage", "View atoms using sage.") 

212define_viewer("x3d", "View atoms using x3d.") 

213 

214# CLI viweers that are internally supported 

215define_viewer( 

216 "avogadro", "View atoms using avogradro.", cli=True, fmt="cube", 

217 argv=["avogadro"] 

218) 

219define_viewer( 

220 "ase_gui_cli", "View atoms using ase gui.", cli=True, fmt="traj", 

221 argv=[sys.executable, '-m', 'ase.gui'], 

222) 

223define_viewer( 

224 "gopenmol", 

225 "View atoms using gopenmol.", 

226 cli=True, 

227 fmt="extxyz", 

228 argv=["runGOpenMol"], 

229) 

230define_viewer( 

231 "rasmol", 

232 "View atoms using rasmol.", 

233 cli=True, 

234 fmt="proteindatabank", 

235 argv=["rasmol", "-pdb"], 

236) 

237define_viewer("vmd", "View atoms using vmd.", cli=True, fmt="cube", 

238 argv=["vmd"]) 

239define_viewer( 

240 "xmakemol", 

241 "View atoms using xmakemol.", 

242 cli=True, 

243 fmt="extxyz", 

244 argv=["xmakemol", "-f"], 

245) 

246 

247register_external_viewer_formats("ase.visualize") 

248 

249CLI_VIEWERS = {key: value for key, value in VIEWERS.items() 

250 if isinstance(value, CLIViewer)} 

251PY_VIEWERS = {key: value for key, value in VIEWERS.items() 

252 if isinstance(value, PyViewer)} 

253 

254 

255def cli_main(): 

256 """ 

257 This is mainly to facilitate launching CLI viewer in a separate python 

258 process 

259 """ 

260 cli_viewer, atoms, data = pickle.load(sys.stdin.buffer) 

261 cli_viewer.view_blocking(atoms, data) 

262 

263 

264if __name__ == "__main__": 

265 cli_main()