|
77 | 77 | from hera.workflows.steps import Step
|
78 | 78 | from hera.workflows.task import Task
|
79 | 79 |
|
80 |
| - |
81 | 80 | _yaml: Optional[ModuleType] = None
|
82 | 81 | try:
|
83 | 82 | import yaml
|
@@ -176,12 +175,28 @@ class ContainerMixin(BaseMixin):
|
176 | 175 | termination_message_policy: Optional[TerminationMessagePolicy] = None
|
177 | 176 | tty: Optional[bool] = None
|
178 | 177 |
|
179 |
| - def _build_image_pull_policy(self) -> Optional[str]: |
| 178 | + def _build_image_pull_policy(self) -> Optional[ImagePullPolicy]: |
180 | 179 | if self.image_pull_policy is None:
|
181 | 180 | return None
|
182 | 181 | 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 |
185 | 200 |
|
186 | 201 | @validator("image", pre=True, always=True)
|
187 | 202 | def _set_image(cls, v):
|
|
0 commit comments