Skip to content

Commit a441f5f

Browse files
author
Ao Guo
committed
Add primitive_or_expr() back to conditions
1 parent 284ddbe commit a441f5f

File tree

1 file changed

+17
-1
lines changed

1 file changed

+17
-1
lines changed

src/sagemaker/workflow/conditions.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,15 @@
2020
import abc
2121

2222
from enum import Enum
23-
from typing import List, Union
23+
from typing import Dict, List, Union
2424

2525
import attr
2626

27+
from sagemaker.workflow import is_pipeline_variable
2728
from sagemaker.workflow.entities import (
2829
DefaultEnumMeta,
2930
Entity,
31+
Expression,
3032
PrimitiveType,
3133
RequestType,
3234
)
@@ -289,3 +291,17 @@ def _referenced_steps(self) -> List[str]:
289291
for condition in self.conditions:
290292
steps.extend(condition._referenced_steps)
291293
return steps
294+
295+
296+
def primitive_or_expr(
297+
value: Union[ExecutionVariable, Expression, PrimitiveType, Parameter, Properties]
298+
) -> Union[Dict[str, str], PrimitiveType]:
299+
"""Provide the expression of the value or return value if it is a primitive.
300+
Args:
301+
value (Union[ConditionValueType, PrimitiveType]): The value to evaluate.
302+
Returns:
303+
Either the expression of the value or the primitive value.
304+
"""
305+
if is_pipeline_variable(value):
306+
return value.expr
307+
return value

0 commit comments

Comments
 (0)