Coverage for /builds/debichem-team/python-ase/ase/gui/save.py: 79.59%

1"""Dialog for saving one or more configurations.""" 

2 

3import numpy as np 

4 

5import ase.gui.ui as ui 

6from ase.gui.i18n import _ 

7from ase.io.formats import ( 

8 filetype, 

9 get_ioformat, 

10 parse_filename, 

11 string2index, 

12 write, 

13) 

14 

15text = _("""\ 

16Append name with "@n" in order to write image 

17number "n" instead of the current image. Append 

18"@start:stop" or "@start:stop:step" if you want 

19to write a range of images. You can leave out 

20"start" and "stop" so that "name@:" will give 

21you all images. Negative numbers count from the 

22last image. Examples: "name@-1": last image, 

23"name@-2:": last two.""") 

24 

25 

26def save_dialog(gui, filename=None): 

27 dialog = ui.SaveFileDialog(gui.window.win, _('Save ...')) 

28 # fix tkinter not automatically setting dialog type 

29 # remove from Python3.8+ 

30 # see https://github.com/python/cpython/pull/25187 

31 # and https://bugs.python.org/issue43655 

32 # and https://github.com/python/cpython/pull/25592 

33 ui.set_windowtype(dialog.top, 'dialog') 

34 ui.Text(text).pack(dialog.top) 

35 filename = filename or dialog.go() 

36 if not filename: 

37 return 

38 

39 filename, index = parse_filename(filename) 

40 if index is None: 

41 index = slice(gui.frame, gui.frame + 1) 

42 elif isinstance(index, str): 

43 index = string2index(index) 

44 elif isinstance(index, slice): 

45 pass 

46 else: 

47 if index < 0: 

48 index += len(gui.images) 

49 index = slice(index, index + 1) 

50 format = filetype(filename, read=False) 

51 io = get_ioformat(format) 

52 

53 extra = {} 

54 remove_hidden = False 

55 if format in ['png', 'eps', 'pov']: 

56 bbox = np.empty(4) 

57 size = gui.window.size / gui.scale 

58 bbox[0:2] = np.dot(gui.center, gui.axes[:, :2]) - size / 2 

59 bbox[2:] = bbox[:2] + size 

60 extra['rotation'] = gui.axes 

61 extra['show_unit_cell'] = gui.window['toggle-show-unit-cell'] 

62 extra['bbox'] = bbox 

63 colors = gui.get_colors(rgb=True) 

64 extra['colors'] = [rgb for rgb, visible 

65 in zip(colors, gui.images.visible) 

66 if visible] 

67 remove_hidden = True 

68 

69 images = [gui.images.get_atoms(i, remove_hidden=remove_hidden) 

70 for i in range(*index.indices(len(gui.images)))] 

71 

72 if len(images) > 1 and io.single: 

73 # We want to write multiple images, but the file format does not 

74 # support it. The solution is to write multiple files, inserting 

75 # a number in the file name before the suffix. 

76 j = filename.rfind('.') 

77 filename = filename[:j] + '{0:05d}' + filename[j:] 

78 for i, atoms in enumerate(images): 

79 write(filename.format(i), atoms, **extra) 

80 else: 

81 try: 

82 write(filename, images, **extra) 

83 except Exception as err: 

84 from ase.gui.ui import showerror 

85 showerror(_('Error'), err) 

86 raise