1import os
2import os.path
3import re
4import shlex
5import shutil
6import subprocess
7
8
9TESTS_DIR = os.path.dirname(__file__)
10TOOL_ROOT = os.path.dirname(TESTS_DIR)
11SRCDIR = os.path.dirname(os.path.dirname(TOOL_ROOT))
12
13MAKE = shutil.which('make')
14FREEZE = os.path.join(TOOL_ROOT, 'freeze.py')
15OUTDIR = os.path.join(TESTS_DIR, 'outdir')
16
17
18class UnsupportedError(Exception):
19    """The operation isn't supported."""
20
21
22def _run_quiet(cmd, cwd=None):
23    #print(f'# {" ".join(shlex.quote(a) for a in cmd)}')
24    try:
25        return subprocess.run(
26            cmd,
27            cwd=cwd,
28            capture_output=True,
29            text=True,
30            check=True,
31        )
32    except subprocess.CalledProcessError as err:
33        # Don't be quiet if things fail
34        print(f"{err.__class__.__name__}: {err}")
35        print("--- STDOUT ---")
36        print(err.stdout)
37        print("--- STDERR ---")
38        print(err.stderr)
39        print("---- END ----")
40        raise
41
42
43def _run_stdout(cmd, cwd=None):
44    proc = _run_quiet(cmd, cwd)
45    return proc.stdout.strip()
46
47
48def find_opt(args, name):
49    opt = f'--{name}'
50    optstart = f'{opt}='
51    for i, arg in enumerate(args):
52        if arg == opt or arg.startswith(optstart):
53            return i
54    return -1
55
56
57def ensure_opt(args, name, value):
58    opt = f'--{name}'
59    pos = find_opt(args, name)
60    if value is None:
61        if pos < 0:
62            args.append(opt)
63        else:
64            args[pos] = opt
65    elif pos < 0:
66        args.extend([opt, value])
67    else:
68        arg = args[pos]
69        if arg == opt:
70            if pos == len(args) - 1:
71                raise NotImplementedError((args, opt))
72            args[pos + 1] = value
73        else:
74            args[pos] = f'{opt}={value}'
75
76
77def copy_source_tree(newroot, oldroot):
78    print(f'copying the source tree into {newroot}...')
79    if os.path.exists(newroot):
80        if newroot == SRCDIR:
81            raise Exception('this probably isn\'t what you wanted')
82        shutil.rmtree(newroot)
83
84    def ignore_non_src(src, names):
85        """Turns what could be a 1000M copy into a 100M copy."""
86        # Don't copy the ~600M+ of needless git repo metadata.
87        # source only, ignore cached .pyc files.
88        subdirs_to_skip = {'.git', '__pycache__'}
89        if os.path.basename(src) == 'Doc':
90            # Another potential ~250M+ of non test related data.
91            subdirs_to_skip.add('build')
92            subdirs_to_skip.add('venv')
93        return subdirs_to_skip
94
95    shutil.copytree(oldroot, newroot, ignore=ignore_non_src)
96    if os.path.exists(os.path.join(newroot, 'Makefile')):
97        _run_quiet([MAKE, 'clean'], newroot)
98
99
100def get_makefile_var(builddir, name):
101    regex = re.compile(rf'^{name} *=\s*(.*?)\s*$')
102    filename = os.path.join(builddir, 'Makefile')
103    try:
104        infile = open(filename, encoding='utf-8')
105    except FileNotFoundError:
106        return None
107    with infile:
108        for line in infile:
109            m = regex.match(line)
110            if m:
111                value, = m.groups()
112                return value or ''
113    return None
114
115
116def get_config_var(builddir, name):
117    python = os.path.join(builddir, 'python')
118    if os.path.isfile(python):
119        cmd = [python, '-c',
120               f'import sysconfig; print(sysconfig.get_config_var("{name}"))']
121        try:
122            return _run_stdout(cmd)
123        except subprocess.CalledProcessError:
124            pass
125    return get_makefile_var(builddir, name)
126
127
128##################################
129# freezing
130
131def prepare(script=None, outdir=None):
132    if not outdir:
133        outdir = OUTDIR
134    os.makedirs(outdir, exist_ok=True)
135
136    # Write the script to disk.
137    if script:
138        scriptfile = os.path.join(outdir, 'app.py')
139        print(f'creating the script to be frozen at {scriptfile}')
140        with open(scriptfile, 'w', encoding='utf-8') as outfile:
141            outfile.write(script)
142
143    # Make a copy of the repo to avoid affecting the current build
144    # (e.g. changing PREFIX).
145    srcdir = os.path.join(outdir, 'cpython')
146    copy_source_tree(srcdir, SRCDIR)
147
148    # We use an out-of-tree build (instead of srcdir).
149    builddir = os.path.join(outdir, 'python-build')
150    os.makedirs(builddir, exist_ok=True)
151
152    # Run configure.
153    print(f'configuring python in {builddir}...')
154    cmd = [
155        os.path.join(srcdir, 'configure'),
156        *shlex.split(get_config_var(srcdir, 'CONFIG_ARGS') or ''),
157    ]
158    ensure_opt(cmd, 'cache-file', os.path.join(outdir, 'python-config.cache'))
159    prefix = os.path.join(outdir, 'python-installation')
160    ensure_opt(cmd, 'prefix', prefix)
161    _run_quiet(cmd, builddir)
162
163    if not MAKE:
164        raise UnsupportedError('make')
165
166    cores = os.cpu_count()
167    if cores and cores >= 3:
168        # this test is most often run as part of the whole suite with a lot
169        # of other tests running in parallel, from 1-2 vCPU systems up to
170        # people's NNN core beasts. Don't attempt to use it all.
171        parallel = f'-j{cores*2//3}'
172    else:
173        parallel = '-j2'
174
175    # Build python.
176    print(f'building python {parallel=} in {builddir}...')
177    if os.path.exists(os.path.join(srcdir, 'Makefile')):
178        # Out-of-tree builds require a clean srcdir.
179        _run_quiet([MAKE, '-C', srcdir, 'clean'])
180    _run_quiet([MAKE, '-C', builddir, parallel])
181
182    # Install the build.
183    print(f'installing python into {prefix}...')
184    _run_quiet([MAKE, '-C', builddir, 'install'])
185    python = os.path.join(prefix, 'bin', 'python3')
186
187    return outdir, scriptfile, python
188
189
190def freeze(python, scriptfile, outdir):
191    if not MAKE:
192        raise UnsupportedError('make')
193
194    print(f'freezing {scriptfile}...')
195    os.makedirs(outdir, exist_ok=True)
196    # Use -E to ignore PYTHONSAFEPATH
197    _run_quiet([python, '-E', FREEZE, '-o', outdir, scriptfile], outdir)
198    _run_quiet([MAKE, '-C', os.path.dirname(scriptfile)])
199
200    name = os.path.basename(scriptfile).rpartition('.')[0]
201    executable = os.path.join(outdir, name)
202    return executable
203
204
205def run(executable):
206    return _run_stdout([executable])
207