1#!/usr/bin/env python3

2#

3# This file is part of Checkbox.

4#

5# Copyright 2022 Canonical Ltd.

6# Written by:

7# Sylvain Pineau <sylvain.pineau@canonical.com>

8#

9# Checkbox is free software: you can redistribute it and/or modify

10# it under the terms of the GNU General Public License version 3,

11# as published by the Free Software Foundation.

12#

13# Checkbox is distributed in the hope that it will be useful,

14# but WITHOUT ANY WARRANTY; without even the implied warranty of

15# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

16# GNU General Public License for more details.

17#

18# You should have received a copy of the GNU General Public License

19# along with Checkbox. If not, see <http://www.gnu.org/licenses/>.

20

21 import argparse

22 import os

23 import subprocess

24 import sys

25 import tarfile

26 import time

27

28 from collections import OrderedDict

29 from fractions import Fraction

30

31 parser = argparse.ArgumentParser()

32 parser.add_argument(

33 "--keyword",

34 default="",

35 help=(

36 "A keyword to distinguish the screenshots "

37 "taken in this run of the script"

38 ),

39)

40 parser.add_argument(

41 "--screenshot-dir",

42 default=os.environ["HOME"],

43 help=(

44 "Specify a directory to store screenshots in. "

45 "Default is %(default)s"

46 ),

47)

48 args = parser.parse_args()

49

50 randrinfo = subprocess.Popen("gnome-randr", shell=True, stdout=subprocess.PIPE)

51 output = randrinfo.communicate()[0].decode().split("\n")

52

53 monitor = ""

54 monitors = dict()

55 highest_modes = [] # list of highest-res modes for each aspect ratio

56 current_modes = [] # remember the user's current settings for cleanup later

57 failures = 0 # count the number of failed modesets

58

59 for line in output:

60 # Ignore Interlaced modes that are indicated by presence of a

61 # trailing 'i' character.

62 if ":" in line or line == "" or "i@" in line:

63 continue

64 if not (line.startswith(" ") or line.startswith("\t")):

65 try:

66 monitor = line.split()[0]

67 monitors[monitor] = OrderedDict()

68 continue

69 except IndexError:

70 continue

71 if monitor:

72 modeline = line.split()

73 try:

74 mode, resolution, rate = modeline[:3]

75 width, height = [int(x) for x in resolution.split("x")]

76 aspect = Fraction(width, height)

77 if width < 675 or width / aspect < 530:

78 continue

79 if resolution in monitors[monitor]:

80 existing_rate = monitors[monitor][resolution][4]

81 if rate < existing_rate:

82 continue

83 monitors[monitor][resolution] = (width, aspect, mode, rate)

84 except IndexError:

85 continue

86

87 for monitor in monitors.keys():

88 # let's create a dict of aspect_ratio:largest_width for each display

89 # (width, because it's easier to compare simple ints when looking for the

90 # highest value).

91 top_res_per_aspect = OrderedDict()

92 connected = False

93 for resolution in monitors[monitor]:

94 width, aspect, mode, rate = monitors[monitor][resolution]

95 cur_max = top_res_per_aspect.get(aspect, 0)

96 top_res_per_aspect[aspect] = max(cur_max, width)

97 if "*" in rate:

98 connected = True

99 current_modes.append((monitor, resolution, mode, rate))

100 if not connected:

101 continue

102 for aspect_ratio, max_width in reversed(top_res_per_aspect.items()):

103 for resolution in monitors[monitor]:

104 width, aspect, mode, rate = monitors[monitor][resolution]

105 if aspect == aspect_ratio and width == max_width:

106 highest_modes.append((monitor, resolution, mode, rate))

107

108 screenshot_path = os.path.join(args.screenshot_dir, "xrandr_screens")

109

110 if args.keyword:

111 screenshot_path = screenshot_path + "_" + args.keyword

112 os.makedirs(screenshot_path, exist_ok=True)

113

114 for monitor, resolution, mode, rate in highest_modes + current_modes:

115 rate = rate.replace("+", "").replace("*", "")

116 print(

117 "Set mode {}@{} for output {}".format(resolution, rate, monitor),

118 flush=True,

119 )

120 cmd = "gnome-randr modify " + monitor + " -m " + mode

121 try:

122 subprocess.run(cmd, check=True, shell=True, stdout=subprocess.PIPE)

123 mode_string = monitor + "_" + resolution

124 filename = os.path.join(screenshot_path, mode_string + ".jpg")

125 cmd = "gnome-screenshot -f " + filename

126 result = subprocess.run(cmd, shell=True, check=False)

127 if result.returncode != 0:

128 print(

129 "Could not capture screenshot -\n"

130 "you may need to install the package 'gnome-screenshot'.",

131 file=sys.stderr,

132 flush=True,

133 )

134 except subprocess.CalledProcessError:

135 failures = failures + 1

136 print(

137 "Failed to set mode {} for output {}:".format(mode, monitor),

138 file=sys.stderr,

139 flush=True,

140 )

141 print(" {}".format(cmd), file=sys.stderr, flush=True)

142 time.sleep(8) # let the hardware recover a bit

143

144# Tar up the screenshots for uploading

145 try:

146 with tarfile.open(screenshot_path + ".tgz", "w:gz") as screen_tar:

147 for screen in os.listdir(screenshot_path):

148 screen_tar.add(screenshot_path + "/" + screen, screen)

149 except (IOError, OSError):

150 pass

151

152 if failures != 0:

153 exit(1)

154 else:

155 exit(0)