From 53395359cd2ff7fa8a0cbeddae7f117093837516 Mon Sep 17 00:00:00 2001
From: Boming Zhang <bomingzh@sjtu.edu.cn>
Date: Sun, 22 Jun 2025 05:55:09 -0400
Subject: [PATCH] feat: files.import-map

---
 joj3_config_generator/models/task.py       |  3 +++
 joj3_config_generator/transformers/task.py | 23 +++++++++++++---------
 tests/convert/basic/task.json              |  3 +++
 tests/convert/basic/task.toml              |  2 +-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/joj3_config_generator/models/task.py b/joj3_config_generator/models/task.py
index 6fca240..6c198dd 100644
--- a/joj3_config_generator/models/task.py
+++ b/joj3_config_generator/models/task.py
@@ -135,6 +135,9 @@ class ParserDiff(BaseModel):
 class StageFiles(BaseModel):
     import_: List[str] = Field([], validation_alias="import")
     export: List[str] = []
+    import_map: Dict[str, str] = Field(
+        {}, validation_alias=AliasChoices("import-map", "import_map")
+    )
 
 
 class Limit(BaseModel):
diff --git a/joj3_config_generator/transformers/task.py b/joj3_config_generator/transformers/task.py
index b3530a8..6b2e92e 100644
--- a/joj3_config_generator/transformers/task.py
+++ b/joj3_config_generator/transformers/task.py
@@ -2,7 +2,7 @@ import re
 import shlex
 from functools import partial
 from pathlib import Path, PurePosixPath
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from joj3_config_generator.models import result, task
 from joj3_config_generator.models.common import Memory, Time
@@ -79,20 +79,25 @@ def get_parser_handler_map(
 def get_executor_with(
     task_stage: task.Stage, cached: Dict[str, None]
 ) -> result.ExecutorWith:
-    file_import = task_stage.files.import_
-    copy_in_files = (file for file in file_import if file not in cached)
+    copy_in: dict[
+        str,
+        Union[result.LocalFile, result.MemoryFile, result.PreparedFile, result.Symlink],
+    ] = {
+        file: result.LocalFile(src=str(JOJ3_CONFIG_ROOT / file))
+        for file in task_stage.files.import_
+        if file not in cached
+    }
+    for src, dst in task_stage.files.import_map.items():
+        if dst in cached:
+            continue
+        copy_in[dst] = result.LocalFile(src=str(JOJ3_CONFIG_ROOT / src))
     file_export = task_stage.files.export
     copy_out_files = ["stdout", "stderr"]
     executor_with_config = result.ExecutorWith(
         default=result.Cmd(
             args=shlex.split(task_stage.command),
             env=[DEFAULT_PATH_ENV, *task_stage.env],
-            copy_in={
-                file: result.LocalFile(src=str(JOJ3_CONFIG_ROOT / file))
-                # all copyin files store in this tools folder
-                # TODO: are there any corner cases?
-                for file in copy_in_files
-            },
+            copy_in=copy_in,
             copy_in_dir="." if task_stage.copy_in_cwd else "",
             copy_out=copy_out_files,
             copy_in_cached={file: file for file in cached},
diff --git a/tests/convert/basic/task.json b/tests/convert/basic/task.json
index 3423db1..64f8f42 100644
--- a/tests/convert/basic/task.json
+++ b/tests/convert/basic/task.json
@@ -143,6 +143,9 @@
                             "copyIn": {
                                 "tools/compile": {
                                     "src": "/home/tt/.config/joj/tools/compile"
+                                },
+                                "h7/Makefile": {
+                                    "src": "/home/tt/.config/joj/tools/Makefile"
                                 }
                             },
                             "copyInCached": {},
diff --git a/tests/convert/basic/task.toml b/tests/convert/basic/task.toml
index cfbdacd..6ceddea 100644
--- a/tests/convert/basic/task.toml
+++ b/tests/convert/basic/task.toml
@@ -20,8 +20,8 @@ name = "Compilation"
 env = [ "CC=clang", "CXX=clang++" ]
 command = "./tools/compile" # eg. script running cmake commands
 files.import = [ "tools/compile" ]
+files.import-map = { "tools/Makefile" = "h7/Makefile" }
 files.export = [ "h7/build/ex2", "h7/build/ex2-asan", "h7/build/ex2-ubsan", "h7/build/ex2-msan", "h7/build/compile_commands.json" ]
-score = 1
 
 # compile parsers
 parsers = [ "result-detail", "result-status" ]