Skip to content

Commit b10dfa5

Browse files
authored
[#687] Adjust image pull policy inference (#689)
**Pull Request Checklist** - [x] Fixes #687 - [ ] Tests added - [x] Documentation/examples added - [x] [Good commit messages](https://cbea.ms/git-commit/) and/or PR title **Description of PR** This PR adds a small mapping of possible user inputs to map to `ImagePullPolicy` objects as a QOL improvement/bug fix --------- Signed-off-by: Flaviu Vadan <[email protected]>
1 parent 854a612 commit b10dfa5

File tree

2 files changed

+23
-5
lines changed

2 files changed

+23
-5
lines changed

src/hera/workflows/_mixins.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@
7777
from hera.workflows.steps import Step
7878
from hera.workflows.task import Task
7979

80-
8180
_yaml: Optional[ModuleType] = None
8281
try:
8382
import yaml
@@ -176,12 +175,28 @@ class ContainerMixin(BaseMixin):
176175
termination_message_policy: Optional[TerminationMessagePolicy] = None
177176
tty: Optional[bool] = None
178177

179-
def _build_image_pull_policy(self) -> Optional[str]:
178+
def _build_image_pull_policy(self) -> Optional[ImagePullPolicy]:
180179
if self.image_pull_policy is None:
181180
return None
182181
elif isinstance(self.image_pull_policy, ImagePullPolicy):
183-
return self.image_pull_policy.value
184-
return ImagePullPolicy[self.image_pull_policy.lower()].value
182+
return self.image_pull_policy
183+
184+
# this helps map image pull policy values as a convenience
185+
policy_mapper = {
186+
# the following 2 are "normal" entries
187+
**{ipp.name: ipp for ipp in ImagePullPolicy},
188+
**{ipp.value: ipp for ipp in ImagePullPolicy},
189+
# some users might submit the policy without underscores
190+
**{ipp.value.lower().replace("_", ""): ipp for ipp in ImagePullPolicy},
191+
# some users might submit the policy in lowercase
192+
**{ipp.name.lower(): ipp for ipp in ImagePullPolicy},
193+
}
194+
try:
195+
return ImagePullPolicy[policy_mapper[self.image_pull_policy].name]
196+
except KeyError as e:
197+
raise KeyError(
198+
f"Supplied image policy {self.image_pull_policy} is not valid. Use one of {ImagePullPolicy.__members__}"
199+
) from e
185200

186201
@validator("image", pre=True, always=True)
187202
def _set_image(cls, v):

src/hera/workflows/script.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,8 @@ def _build_template(self) -> _ModelTemplate:
194194

195195
def _build_script(self) -> _ModelScriptTemplate:
196196
assert isinstance(self.constructor, ScriptConstructor)
197+
image_pull_policy = self._build_image_pull_policy()
198+
197199
return self.constructor.transform_script_template_post_build(
198200
self,
199201
_ModelScriptTemplate(
@@ -202,7 +204,8 @@ def _build_script(self) -> _ModelScriptTemplate:
202204
env=self._build_env(),
203205
env_from=self._build_env_from(),
204206
image=self.image,
205-
image_pull_policy=self._build_image_pull_policy(),
207+
# `image_pull_policy` in script wants a string not an `ImagePullPolicy` object
208+
image_pull_policy=None if image_pull_policy is None else image_pull_policy.value,
206209
lifecycle=self.lifecycle,
207210
liveness_probe=self.liveness_probe,
208211
name=self.container_name,

0 commit comments

Comments
 (0)