get_files_to_run shoudl handle extra files

Add option (to be used by #2006) to specify extra files particular tutorial depends on
Make sure those files will not get deleted

Add a bit more type annotation to the code

Author: malfet

.jenkins/get_files_to_run.py

@@ -1,4 +1,4 @@
-from typing import List, Any
+from typing import Any, Dict, List, Optional, Tuple
 import json
 import os
 from pathlib import Path
@@ -14,17 +14,21 @@ def get_all_files() -> List[str]:
     return [str(x) for x in sources]
 
 
-def calculate_shards(all_files: List[str], num_shards: int = 20):
+def read_metadata() -> Dict[str, Any]:
     with (REPO_BASE_DIR / ".jenkins" / "metadata.json").open() as fp:
-        metadata = json.load(fp)
-    sharded_files = [(0.0, []) for _ in range(num_shards)]
+        return json.load(fp)
 
-    def get_duration(file):
+
+def calculate_shards(all_files: List[str], num_shards: int = 20) -> List[List[str]]:
+    sharded_files: List[Tuple[float, List[str]]] = [(0.0, []) for _ in range(num_shards)]
+    metadata = read_metadata()
+
+    def get_duration(file: str) -> int:
         # tutorials not listed in the metadata.json file usually take
         # <3min to run, so we'll default to 1min if it's not listed
         return metadata.get(file, {}).get("duration", 60)
 
-    def get_needs_machine(file):
+    def get_needs_machine(file: str) -> Optional[str]:
         return metadata.get(file, {}).get("needs", None)
 
     def add_to_shard(i, filename):
@@ -55,9 +59,19 @@ def add_to_shard(i, filename):
     return [x[1] for x in sharded_files]
 
 
-def remove_other_files(all_files, files_to_run) -> None:
+def compute_files_to_keep(files_to_run: List[str]) -> List[str]:
+    metadata = read_metadata()
+    files_to_keep = list(files_to_run)
+    for file in files_to_run:
+        extra_files = metadata.get(file, {}).get("extra_files", [])
+        files_to_keep.extend(extra_files)
+    return files_to_keep
+
+
+def remove_other_files(all_files, files_to_keep) -> None:
+
     for file in all_files:
-        if file not in files_to_run:
+        if file not in files_to_keep:
             remove_runnable_code(file, file)
 
 
@@ -76,7 +90,7 @@ def main() -> None:
     all_files = get_all_files()
     files_to_run = calculate_shards(all_files, num_shards=args.num_shards)[args.shard_num]
     if not args.dry_run:
-        remove_other_files(all_files, files_to_run)
+        remove_other_files(all_files, compute_files_to_keep(files_to_run))
     stripped_file_names = [Path(x).stem for x in files_to_run]
     print(" ".join(stripped_file_names))