recompile_requirements: Add "#@ pre" marker

1 files changed, 23 insertions, 10 deletions

diff --git a/scripts/dev/recompile_requirements.py b/scripts/dev/recompile_requirements.py
index ecd7dd153..3e7db0c9a 100644
--- a/scripts/dev/recompile_requirements.py
+++ b/scripts/dev/recompile_requirements.py
@@ -81,12 +81,18 @@ def read_comments(fobj):
         'ignore': [],
         'add': [],
         'replace': {},
+        'pre': False,
     }
     for line in fobj:
         if line.startswith('#@'):
-            command, args = line[2:].split(':', maxsplit=1)
-            command = command.strip()
-            args = args.strip()
+            if ':' in line:
+                command, args = line[2:].split(':', maxsplit=1)
+                command = command.strip()
+                args = args.strip()
+            else:
+                command = line[2:].strip()
+                args = None
+
             if command == 'filter':
                 pkg, filt = args.split(' ', maxsplit=1)
                 comments['filter'][pkg] = filt
@@ -103,6 +109,8 @@ def read_comments(fobj):
                 comments['markers'][pkg] = markers
             elif command == 'add':
                 comments['add'].append(args)
+            elif command == 'pre':
+                comments['pre'] = True
     return comments
 
 
@@ -113,7 +121,7 @@ def get_all_names():
         yield basename[len('requirements-'):-len('.txt-raw')]
 
 
-def init_venv(host_python, venv_dir, requirements):
+def init_venv(host_python, venv_dir, requirements, pre=False):
     """Initialize a new virtualenv and install the given packages."""
     subprocess.run([host_python, '-m', 'venv', venv_dir], check=True)
 
@@ -121,8 +129,10 @@ def init_venv(host_python, venv_dir, requirements):
     subprocess.run([venv_python, '-m', 'pip',
                     'install', '-U', 'pip'], check=True)
 
-    subprocess.run([venv_python, '-m', 'pip',
-                    'install', '-r', requirements], check=True)
+    install_command = [venv_python, '-m', 'pip', 'install', '-r', requirements]
+    if pre:
+        install_command.append('--pre')
+    subprocess.run(install_command, check=True)
     subprocess.run([venv_python, '-m', 'pip', 'check'], check=True)
     return venv_python
 
@@ -152,15 +162,18 @@ def main():
 
         utils.print_subtitle("Building")
 
+        with open(filename, 'r', encoding='utf-8') as f:
+            comments = read_comments(f)
+
         with tempfile.TemporaryDirectory() as tmpdir:
-            venv_python = init_venv(host_python, tmpdir, filename)
+            venv_python = init_venv(host_python=host_python,
+                                    venv_dir=tmpdir,
+                                    requirements=filename,
+                                    pre=comments['pre'])
             proc = subprocess.run([venv_python, '-m', 'pip', 'freeze'],
                                   check=True, stdout=subprocess.PIPE)
             reqs = proc.stdout.decode('utf-8')
 
-        with open(filename, 'r', encoding='utf-8') as f:
-            comments = read_comments(f)
-
         with open(outfile, 'w', encoding='utf-8') as f:
             f.write("# This file is automatically generated by "
                     "scripts/dev/recompile_requirements.py\n\n")