1#!/usr/bin/env python3
2
3#
4# Copyright 2024, The Android Open Source Project
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18"""Script to update golden screenshot files from a temp directory with all files."""
19
20import argparse
21import os
22import pathlib
23import re
24import sys
25import shutil
26
27ACTUAL_SCREENSHOT_FILE_LINE_PATTERN = r'image_location_test: "(?P<name>[^"]+)"'
28GOLDEN_PATH_LINE_PATTERN = r'image_location_golden: "(?P<name>[^"]+)"'
29
30def main():
31  parser = argparse.ArgumentParser(
32      'Update golden screenshots files in a temporary directory.')
33  parser.add_argument(
34      '--android-build-top',
35      default='main',
36      help='The Android build top directory.')
37  parser.add_argument(
38      '--source-directory',
39      default='',
40      help='The directory where all proto buffer and screenshot files are located.')
41
42  args = parser.parse_args()
43
44  if not args.source_directory:
45    print('The directory which contain all proto buffer and screenshot files must be specified.')
46    sys.exit(1)
47
48  source_directory = os.path.normpath(args.source_directory)
49
50  if not os.path.exists(source_directory):
51    print('The source directory does not exist: ' + source_directory)
52    sys.exit(1)
53
54  if not os.path.isdir(source_directory):
55    print('The specified path is not a directory: ' + source_directory)
56    sys.exit(1)
57
58  # Due to file renaming, a map from the desired actual screenshot file name
59  # needs to be mapped to the file name after renaming should be pre-built.
60  actual_screenshot_files = {}
61  for filename in os.listdir(args.source_directory):
62    if not (filename.endswith('.png') or filename.find('_actual') < 0):
63      continue
64    actual_screenshot_files[get_stripped_actual_screenshot_file(filename)] = filename
65
66  for filename in os.listdir(args.source_directory):
67    if not (filename.endswith('.txt') or filename.endswith('.textpb')):
68      continue
69
70    pb_file = open(os.path.join(args.source_directory, filename), 'r')
71    actual_screenshot_file_name = ''
72    golden_path = ''
73
74    # Parse all text protos which contain information of golden path and actual
75    # screenshot filenames.
76    for lines in pb_file.readlines():
77      match_actual = re.search(ACTUAL_SCREENSHOT_FILE_LINE_PATTERN, lines)
78      if match_actual:
79        actual_screenshot_file_name = match_actual.group('name')
80        continue
81      match_golden = re.search(GOLDEN_PATH_LINE_PATTERN, lines)
82      if match_golden:
83        golden_path = match_golden.group('name')
84
85    if not golden_path or not actual_screenshot_file_name:
86      print(
87          f'Golden path: {golden_path} or '
88          f'Actual screenshot name: {actual_screenshot_file_name} is empty')
89      continue
90
91    # Copy the actual file to the destination.
92    src_path = os.path.join(
93        args.source_directory,
94        actual_screenshot_files[actual_screenshot_file_name])
95    dest_path = os.path.join(pathlib.Path.home(), args.android_build_top, golden_path)
96    shutil.copyfile(src_path, dest_path)
97    print(f'Updated {dest_path}')
98
99def get_stripped_actual_screenshot_file(original_file_name: str) -> str:
100  first_index = original_file_name.find('.png_')
101  if first_index < 0:
102    return original_file_name
103  else:
104    return original_file_name[:first_index] + '.png'
105
106if __name__ == '__main__':
107  main()
108