1# Copyright 2015 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import argparse
16import glob
17import multiprocessing
18import os
19import pickle
20import shutil
21import sys
22import tempfile
23from typing import Dict, List, Union
24
25import _utils
26import yaml
27
28PROJECT_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..",
29                            "..")
30os.chdir(PROJECT_ROOT)
31# TODO(lidiz) find a better way for plugins to reference each other
32sys.path.append(os.path.join(PROJECT_ROOT, 'tools', 'buildgen', 'plugins'))
33
34# from tools.run_tests.python_utils import jobset
35jobset = _utils.import_python_module(
36    os.path.join(PROJECT_ROOT, 'tools', 'run_tests', 'python_utils',
37                 'jobset.py'))
38
39PREPROCESSED_BUILD = '.preprocessed_build'
40test = {} if os.environ.get('TEST', 'false') == 'true' else None
41
42assert sys.argv[1:], 'run generate_projects.sh instead of this directly'
43parser = argparse.ArgumentParser()
44parser.add_argument('build_files',
45                    nargs='+',
46                    default=[],
47                    help="build files describing build specs")
48parser.add_argument('--templates',
49                    nargs='+',
50                    default=[],
51                    help="mako template files to render")
52parser.add_argument('--output_merged',
53                    '-m',
54                    default='',
55                    type=str,
56                    help="merge intermediate results to a file")
57parser.add_argument('--jobs',
58                    '-j',
59                    default=multiprocessing.cpu_count(),
60                    type=int,
61                    help="maximum parallel jobs")
62parser.add_argument('--base',
63                    default='.',
64                    type=str,
65                    help="base path for generated files")
66args = parser.parse_args()
67
68
69def preprocess_build_files() -> _utils.Bunch:
70    """Merges build yaml into a one dictionary then pass it to plugins."""
71    build_spec = dict()
72    for build_file in args.build_files:
73        with open(build_file, 'r') as f:
74            _utils.merge_json(build_spec, yaml.safe_load(f.read()))
75    # Executes plugins. Plugins update the build spec in-place.
76    for py_file in sorted(glob.glob('tools/buildgen/plugins/*.py')):
77        plugin = _utils.import_python_module(py_file)
78        plugin.mako_plugin(build_spec)
79    if args.output_merged:
80        with open(args.output_merged, 'w') as f:
81            f.write(yaml.dump(build_spec))
82    # Makes build_spec sort of immutable and dot-accessible
83    return _utils.to_bunch(build_spec)
84
85
86def generate_template_render_jobs(templates: List[str]) -> List[jobset.JobSpec]:
87    """Generate JobSpecs for each one of the template rendering work."""
88    jobs = []
89    base_cmd = [sys.executable, 'tools/buildgen/_mako_renderer.py']
90    for template in sorted(templates, reverse=True):
91        root, f = os.path.split(template)
92        if os.path.splitext(f)[1] == '.template':
93            out_dir = args.base + root[len('templates'):]
94            out = os.path.join(out_dir, os.path.splitext(f)[0])
95            if not os.path.exists(out_dir):
96                os.makedirs(out_dir)
97            cmd = base_cmd[:]
98            cmd.append('-P')
99            cmd.append(PREPROCESSED_BUILD)
100            cmd.append('-o')
101            if test is None:
102                cmd.append(out)
103            else:
104                tf = tempfile.mkstemp()
105                test[out] = tf[1]
106                os.close(tf[0])
107                cmd.append(test[out])
108            cmd.append(args.base + '/' + root + '/' + f)
109            jobs.append(jobset.JobSpec(cmd, shortname=out,
110                                       timeout_seconds=None))
111    return jobs
112
113
114def main() -> None:
115    templates = args.templates
116    if not templates:
117        for root, _, files in os.walk('templates'):
118            for f in files:
119                templates.append(os.path.join(root, f))
120
121    build_spec = preprocess_build_files()
122    with open(PREPROCESSED_BUILD, 'wb') as f:
123        pickle.dump(build_spec, f)
124
125    err_cnt, _ = jobset.run(generate_template_render_jobs(templates),
126                            maxjobs=args.jobs)
127    if err_cnt != 0:
128        print('ERROR: %s error(s) found while generating projects.' % err_cnt,
129              file=sys.stderr)
130        sys.exit(1)
131
132    if test is not None:
133        for s, g in test.items():
134            if os.path.isfile(g):
135                assert 0 == os.system('diff %s %s' % (s, g)), s
136                os.unlink(g)
137            else:
138                assert 0 == os.system('diff -r %s %s' % (s, g)), s
139                shutil.rmtree(g, ignore_errors=True)
140
141
142if __name__ == "__main__":
143    main()
144