diff --git a/src/openenv/core/containers/runtime/novita_provider.py b/src/openenv/core/containers/runtime/novita_provider.py index 6b5926efa..f0ed9735e 100644 --- a/src/openenv/core/containers/runtime/novita_provider.py +++ b/src/openenv/core/containers/runtime/novita_provider.py @@ -62,6 +62,7 @@ _ARG_DEFAULT_RE = re.compile( r"^\s*ARG\s+(?P\w+)=(?P\S+)\s*$", re.IGNORECASE ) +_ARG_REFERENCE_RE = re.compile(r"\$(?:\{(?P\w+)\}|(?P\w+))") def _strip_mount_flags(content: str) -> str: @@ -116,6 +117,23 @@ def _strip_mount_flags(content: str) -> str: return "\n".join(out) +def _substitute_args(reference: str, arg_defaults: Dict[str, str]) -> str: + """Expand ``$NAME`` and ``${NAME}`` in one pass against *arg_defaults*. + + Substituting name by name would let a shorter ARG consume the prefix of a + longer one: with ``ARG BASE`` declared before ``ARG BASE_IMAGE``, an + unbraced ``FROM $BASE_IMAGE`` became ``_IMAGE``. Matching the + whole name in a single pass makes the result independent of declaration + order. Names with no default are left verbatim, as before. + """ + + def replace(match: "re.Match[str]") -> str: + name = match.group("braced") or match.group("bare") + return arg_defaults.get(name, match.group(0)) + + return _ARG_REFERENCE_RE.sub(replace, reference) + + def _resolve_from_references(content: str) -> str: """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``. @@ -130,9 +148,8 @@ def _resolve_from_references(content: str) -> str: as part of the name. Novita builds for its own platform, so the flag is dropped rather than propagated. - Substitutions are applied longest-name-first so a short unbraced prefix - cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite - ``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``). + Substitutions match whole names in one pass, leaving undeclared names + intact even when a declared ARG is a prefix of the reference. """ arg_defaults: Dict[str, str] = {} for line in content.split("\n"): @@ -142,11 +159,6 @@ def _resolve_from_references(content: str) -> str: if match: arg_defaults[match.group("name")] = match.group("value") - # Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token. - arg_items = sorted( - arg_defaults.items(), key=lambda item: len(item[0]), reverse=True - ) - out: List[str] = [] for line in content.split("\n"): match = _FROM_RE.match(line) @@ -156,9 +168,7 @@ def _resolve_from_references(content: str) -> str: reference = match.group("rest").strip() reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip() - for name, value in arg_items: - reference = reference.replace(f"${{{name}}}", value) - reference = reference.replace(f"${name}", value) + reference = _substitute_args(reference, arg_defaults) out.append(f"FROM {reference}") return "\n".join(out) diff --git a/tests/test_core/test_novita_provider.py b/tests/test_core/test_novita_provider.py index 21e879da9..bc83847a1 100644 --- a/tests/test_core/test_novita_provider.py +++ b/tests/test_core/test_novita_provider.py @@ -902,6 +902,64 @@ def test_arg_after_first_from_not_in_scope(self): out = _resolve_from_references("FROM python:3.11\nARG X=alpine\nRUN echo hi\n") assert "FROM python:3.11" in out + @pytest.mark.parametrize("reference", ["$BASE_IMAGE", "${BASE_IMAGE}"]) + def test_shorter_arg_does_not_consume_a_longer_name(self, reference): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + # ARG BASE is declared first, so name-by-name substitution used to turn + # `FROM $BASE_IMAGE` into `python:3.12_IMAGE`. + out = _resolve_from_references( + "ARG BASE=python:3.12\n" + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + f"FROM {reference} AS builder\n" + ) + assert "FROM ghcr.io/huggingface/openenv-base:latest AS builder" in out + + def test_substitution_is_independent_of_declaration_order(self): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + declarations = [ + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest", + "ARG BASE=python:3.12", + ] + outputs = { + _resolve_from_references("\n".join([*order, "FROM $BASE_IMAGE\n"])) + for order in (declarations, list(reversed(declarations))) + } + assert outputs == { + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + "ARG BASE=python:3.12\n" + "FROM ghcr.io/huggingface/openenv-base:latest\n", + "ARG BASE=python:3.12\n" + "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n" + "FROM ghcr.io/huggingface/openenv-base:latest\n", + } + + @pytest.mark.parametrize( + "reference", ["$UNSET_IMAGE", "$BASE_IMAGE", "${BASE_IMAGE}"] + ) + def test_undeclared_arg_reference_is_left_verbatim(self, reference): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + out = _resolve_from_references(f"ARG BASE=python:3.12\nFROM {reference}\n") + assert f"FROM {reference}" in out + + def test_substituted_value_is_not_expanded_again(self): + from openenv.core.containers.runtime.novita_provider import ( + _resolve_from_references, + ) + + out = _resolve_from_references( + "ARG REGISTRY=example.com\nARG IMAGE=$REGISTRY/app\nFROM $IMAGE\n" + ) + assert "FROM $REGISTRY/app" in out + def test_flatten_drops_same_path_copy(self): from openenv.core.containers.runtime.novita_provider import _flatten_multistage