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

2 

3from ase.gui.i18n import _ 

4 

5import numpy as np 

6 

7import ase.gui.ui as ui 

8from ase.io.formats import (write, parse_filename, get_ioformat, string2index, 

9 filetype) 

10 

11 

12text = _("""\ 

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

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

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

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

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

18you all images. Negative numbers count from the 

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

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

21 

22 

23def save_dialog(gui, filename=None): 

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

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

26 filename = filename or dialog.go() 

27 if not filename: 

28 return 

29 

30 filename, index = parse_filename(filename) 

31 if index is None: 

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

33 elif isinstance(index, str): 

34 index = string2index(index) 

35 elif isinstance(index, slice): 

36 pass 

37 else: 

38 if index < 0: 

39 index += len(gui.images) 

40 index = slice(index, index + 1) 

41 format = filetype(filename, read=False) 

42 io = get_ioformat(format) 

43 

44 extra = {} 

45 remove_hidden = False 

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

47 bbox = np.empty(4) 

48 size = gui.window.size / gui.scale 

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

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

51 extra['rotation'] = gui.axes 

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

53 extra['bbox'] = bbox 

54 colors = gui.get_colors(rgb=True) 

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

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

57 if visible] 

58 remove_hidden = True 

59 

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

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

62 

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

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

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

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

67 j = filename.rfind('.') 

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

69 for i, atoms in enumerate(images): 

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

71 else: 

72 try: 

73 write(filename, images, **extra) 

74 except Exception as err: 

75 from ase.gui.ui import showerror 

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

77 raise