1from io import BytesIO 

2import pickle 

3import subprocess 

4import sys 

5import tempfile 

6from pathlib import Path 

7from contextlib import contextmanager 

8 

9from ase.io.formats import ioformats 

10from ase.io import write 

11 

12 

13def _pipe_to_ase_gui(atoms, repeat): 

14 buf = BytesIO() 

15 write(buf, atoms, format='traj') 

16 

17 args = [sys.executable, '-m', 'ase', 'gui', '-'] 

18 if repeat: 

19 args.append('--repeat={},{},{}'.format(*repeat)) 

20 

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

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

23 proc.stdin.close() 

24 return proc 

25 

26 

27class CLIViewer: 

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

29 self.name = name 

30 self.fmt = fmt 

31 self.argv = argv 

32 

33 @property 

34 def ioformat(self): 

35 return ioformats[self.fmt] 

36 

37 @contextmanager 

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

39 ioformat = self.ioformat 

40 suffix = '.' + ioformat.extensions[0] 

41 

42 if ioformat.isbinary: 

43 mode = 'wb' 

44 else: 

45 mode = 'w' 

46 

47 with tempfile.TemporaryDirectory(prefix='ase-view-') as dirname: 

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

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

50 # cannot be open on multiple processes). 

51 path = Path(dirname) / f'atoms{suffix}' 

52 with path.open(mode) as fd: 

53 if data is None: 

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

55 else: 

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

57 yield path 

58 

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

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

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

62 

63 def view(self, atoms, data=None, repeat=None): 

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

65 if repeat is not None: 

66 atoms = atoms.repeat(repeat) 

67 

68 proc = subprocess.Popen( 

69 [sys.executable, '-m', 'ase.visualize.external'], 

70 stdin=subprocess.PIPE) 

71 

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

73 proc.stdin.close() 

74 return proc 

75 

76 @classmethod 

77 def viewers(cls): 

78 # paraview_script = Path(__file__).parent / 'paraview_script.py' 

79 # Can we make paraview/vtkxml work on some test system? 

80 return [ 

81 cls('ase_gui_cli', 'traj', [sys.executable, '-m', 'ase.gui']), 

82 cls('avogadro', 'cube', ['avogadro']), 

83 cls('gopenmol', 'extxyz', ['runGOpenMol']), 

84 cls('rasmol', 'proteindatabank', ['rasmol', '-pdb']), 

85 cls('vmd', 'cube', ['vmd']), 

86 cls('xmakemol', 'extxyz', ['xmakemol', '-f']), 

87 # cls('paraview', 'vtu', 

88 # ['paraview', f'--script={paraview_script}']) 

89 ] 

90 

91 

92class PyViewer: 

93 def __init__(self, name, supports_repeat=False): 

94 self.name = name 

95 self.supports_repeat = supports_repeat 

96 

97 def view(self, atoms, data=None, repeat=None): 

98 # Delegate to any of the below methods 

99 func = getattr(self, self.name) 

100 if self.supports_repeat: 

101 return func(atoms, repeat) 

102 else: 

103 if repeat is not None: 

104 atoms = atoms.repeat(repeat) 

105 return func(atoms) 

106 

107 def sage(self, atoms): 

108 from ase.visualize.sage import view_sage_jmol 

109 return view_sage_jmol(atoms) 

110 

111 def ngl(self, atoms): 

112 from ase.visualize.ngl import view_ngl 

113 return view_ngl(atoms) 

114 

115 def x3d(self, atoms): 

116 from ase.visualize.x3d import view_x3d 

117 return view_x3d(atoms) 

118 

119 def ase(self, atoms, repeat): 

120 return _pipe_to_ase_gui(atoms, repeat) 

121 

122 @classmethod 

123 def viewers(cls): 

124 return [ 

125 cls('ase', supports_repeat=True), 

126 cls('ngl'), 

127 cls('sage'), 

128 cls('x3d'), 

129 ] 

130 

131 

132viewers = {viewer.name: viewer 

133 for viewer in CLIViewer.viewers() + PyViewer.viewers()} 

134viewers['nglview'] = viewers['ngl'] 

135 

136 

137def main(): 

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

139 cli_viewer.view_blocking(atoms, data) 

140 

141 

142if __name__ == '__main__': 

143 main()