From 398de4bd23fa3bfd733918a1b9d0bd31ff058da6 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 10 Apr 2026 22:05:13 +0100 Subject: [PATCH 1/5] Add tests Signed-off-by: Thomas Newton --- .../unit/core/test_python_auto_container.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index ed849c34e2..0e4df6ce7b 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -69,6 +69,36 @@ def test_image_name_interpolation(default_image_config): assert img == "flyte/test:d1" +@pytest.mark.parametrize( + "image,expected_resolved_image", + [ + ( + Image(name="image_name", fqn="docker.io/xyz", tag="tag"), + "docker.io/xyz:tag--docker.io/xyz:tag", + ), + ( + Image(name="image-name", fqn="docker.io/xyz", tag="tag"), + "docker.io/xyz:tag--docker.io/xyz:tag", + ), + ( + Image(name="image-name", fqn="docker.io/xyz", digest="sha256:abc123"), + "docker.io/xyz:sha256:abc123--docker.io/xyz@sha256:abc123", + ), + ( + Image(name="image-name", fqn="xyz", tag="tag"), + "xyz:tag--xyz:tag", + ), + ], +) +def test_image_interpolation_with_various_image_names(default_image_config, image, expected_resolved_image): + image_config = ImageConfig.create_from( + default_image=default_image_config.default_image, + other_images=[image], + ) + # This isn't meant to be realistic, but it covers, fqnm version and full image interpolation. + templated_image = f"{{{{.image.{image.name}.fqn}}}}:{{{{.image.{image.name}.version}}}}--{{{{.image.{image.name}}}}}" + assert get_registerable_container_image(img=templated_image, cfg=image_config) == expected_resolved_image + class DummyAutoContainerTask(PythonAutoContainerTask): def execute(self, **kwargs) -> Any: pass From cf811262ace900e6cd11a385c52ab816854613c7 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 10 Apr 2026 22:05:22 +0100 Subject: [PATCH 2/5] Update regex Signed-off-by: Thomas Newton --- flytekit/core/python_auto_container.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index f18f3c1bca..3ae647df60 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -489,4 +489,4 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: # fqn will access the fully qualified name of the image (e.g. registry/imagename:version -> registry/imagename) # version will access the version part of the image (e.g. registry/imagename:version -> version) # With empty attribute, it'll access the full image path (e.g. registry/imagename:version -> registry/imagename:version) -_IMAGE_REPLACE_REGEX = re.compile(r"({{\s*\.image[s]?(?:\.([a-zA-Z0-9_]+))(?:\.([a-zA-Z0-9_]+))?\s*}})", re.IGNORECASE) +_IMAGE_REPLACE_REGEX = re.compile(r"({{\s*\.image[s]?(?:\.([a-zA-Z0-9_-]+))(?:\.(fqn|version))?\s*}})", re.IGNORECASE) From 11af75b82c9817a740a514846ffb117ae69916c1 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 10 Apr 2026 22:14:32 +0100 Subject: [PATCH 3/5] Fix other test Signed-off-by: Thomas Newton --- flytekit/core/python_auto_container.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 3ae647df60..fe0fc37579 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -489,4 +489,5 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: # fqn will access the fully qualified name of the image (e.g. registry/imagename:version -> registry/imagename) # version will access the version part of the image (e.g. registry/imagename:version -> version) # With empty attribute, it'll access the full image path (e.g. registry/imagename:version -> registry/imagename:version) -_IMAGE_REPLACE_REGEX = re.compile(r"({{\s*\.image[s]?(?:\.([a-zA-Z0-9_-]+))(?:\.(fqn|version))?\s*}})", re.IGNORECASE) +_IMAGE_TEMPLATE_SUBSTRING_REGEX = r"(?:\.([a-zA-Z0-9_-]+))" +_IMAGE_REPLACE_REGEX = re.compile(f"({{{{\s*\.image[s]?{_IMAGE_TEMPLATE_SUBSTRING_REGEX}{_IMAGE_TEMPLATE_SUBSTRING_REGEX}?\s*}}}})", re.IGNORECASE) From a3a26d562ea882e2a43524624028f90b3a997845 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 10 Apr 2026 22:22:30 +0100 Subject: [PATCH 4/5] Move tests somewhere more appropriate Signed-off-by: Thomas Newton --- .../unit/core/test_python_auto_container.py | 30 ------------------- .../unit/core/test_python_function_task.py | 21 ++++++++++++- 2 files changed, 20 insertions(+), 31 deletions(-) diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index 0e4df6ce7b..ed849c34e2 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -69,36 +69,6 @@ def test_image_name_interpolation(default_image_config): assert img == "flyte/test:d1" -@pytest.mark.parametrize( - "image,expected_resolved_image", - [ - ( - Image(name="image_name", fqn="docker.io/xyz", tag="tag"), - "docker.io/xyz:tag--docker.io/xyz:tag", - ), - ( - Image(name="image-name", fqn="docker.io/xyz", tag="tag"), - "docker.io/xyz:tag--docker.io/xyz:tag", - ), - ( - Image(name="image-name", fqn="docker.io/xyz", digest="sha256:abc123"), - "docker.io/xyz:sha256:abc123--docker.io/xyz@sha256:abc123", - ), - ( - Image(name="image-name", fqn="xyz", tag="tag"), - "xyz:tag--xyz:tag", - ), - ], -) -def test_image_interpolation_with_various_image_names(default_image_config, image, expected_resolved_image): - image_config = ImageConfig.create_from( - default_image=default_image_config.default_image, - other_images=[image], - ) - # This isn't meant to be realistic, but it covers, fqnm version and full image interpolation. - templated_image = f"{{{{.image.{image.name}.fqn}}}}:{{{{.image.{image.name}.version}}}}--{{{{.image.{image.name}}}}}" - assert get_registerable_container_image(img=templated_image, cfg=image_config) == expected_resolved_image - class DummyAutoContainerTask(PythonAutoContainerTask): def execute(self, **kwargs) -> Any: pass diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index f09e8e38f5..427e3bd774 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -48,7 +48,17 @@ def test_container_image_conversion(mock_image_spec_builder): name="other3", fqn="xyz.com/other3", ) - cfg = ImageConfig(default_image=default_img, images=[default_img, other_img, other_img2, other_img3]) + other_img4 = Image( + name="other-4", + fqn="other4", + tag="tag4" + ) + other_img5 = Image( + name="other_5", + fqn="xyz.com/other5", + tag="tag5" + ) + cfg = ImageConfig(default_image=default_img, images=[default_img, other_img, other_img2, other_img3, other_img4, other_img5]) assert get_registerable_container_image(None, cfg) == "xyz.com/abc:tag1" assert get_registerable_container_image("", cfg) == "xyz.com/abc:tag1" assert get_registerable_container_image("abc", cfg) == "abc" @@ -72,6 +82,15 @@ def test_container_image_conversion(mock_image_spec_builder): get_registerable_container_image("{{.image.other3.fqn}}:{{.image.other3.version}}", cfg) == "xyz.com/other3:tag1" ) + assert ( + get_registerable_container_image("{{.image.other-4.fqn}}:{{.image.other-4.version}}", cfg) + == "other4:tag4" + ) + assert ( + get_registerable_container_image("{{.image.other_5.fqn}}:{{.image.other_5.version}}", cfg) + == "xyz.com/other5:tag5" + ) + assert get_registerable_container_image("{{.image.other.fqn}}", cfg) == "xyz.com/other" # Works with images instead of just image assert get_registerable_container_image("{{.images.other.fqn}}", cfg) == "xyz.com/other" From b028292347889ccc56b46a302a51f0d7259b01aa Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 10 Apr 2026 23:04:42 +0100 Subject: [PATCH 5/5] Autoformat Signed-off-by: Thomas Newton --- flytekit/core/python_auto_container.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index fe0fc37579..0e3dd5bda8 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -490,4 +490,6 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: # version will access the version part of the image (e.g. registry/imagename:version -> version) # With empty attribute, it'll access the full image path (e.g. registry/imagename:version -> registry/imagename:version) _IMAGE_TEMPLATE_SUBSTRING_REGEX = r"(?:\.([a-zA-Z0-9_-]+))" -_IMAGE_REPLACE_REGEX = re.compile(f"({{{{\s*\.image[s]?{_IMAGE_TEMPLATE_SUBSTRING_REGEX}{_IMAGE_TEMPLATE_SUBSTRING_REGEX}?\s*}}}})", re.IGNORECASE) +_IMAGE_REPLACE_REGEX = re.compile( + rf"({{{{\s*\.image[s]?{_IMAGE_TEMPLATE_SUBSTRING_REGEX}{_IMAGE_TEMPLATE_SUBSTRING_REGEX}?\s*}}}})", re.IGNORECASE +)