feature: support JsonGet for all step types

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -6,7 +6,7 @@ ConditionStep
 
 .. autoclass:: sagemaker.workflow.condition_step.ConditionStep
 
-.. autoclass:: sagemaker.workflow.condition_step.JsonGet
+.. deprecated:: sagemaker.workflow.condition_step.JsonGet
 
 Conditions
 ----------
@@ -54,6 +54,8 @@ Functions
 
 .. autoclass:: sagemaker.workflow.functions.Join
 
+.. autoclass:: sagemaker.workflow.functions.JsonGet
+
 Parameters
 ----------
 

src/sagemaker/workflow/condition_step.py

@@ -17,6 +17,7 @@
 
 import attr
 
+from sagemaker.deprecations import deprecated_class
 from sagemaker.workflow.conditions import Condition
 from sagemaker.workflow.entities import (
     Expression,
@@ -114,3 +115,6 @@ def expr(self):
                 "Path": self.json_path,
             }
         }
+
+
+JsonGet = deprecated_class(JsonGet, "JsonGet")

src/sagemaker/workflow/functions.py

@@ -13,11 +13,12 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List
+from typing import List, Union
 
 import attr
 
 from sagemaker.workflow.entities import Expression
+from sagemaker.workflow.properties import PropertyFile
 
 
 @attr.s
@@ -54,3 +55,36 @@ def expr(self):
                 ],
             },
         }
+
+
+@attr.s
+class JsonGet(Expression):
+    """Get JSON properties from PropertyFiles.
+
+    Attributes:
+        step_name (str): The step name from which to get the property file.
+        property_file (Union[PropertyFile, str]): Either a PropertyFile instance
+            or the name of a property file.
+        json_path (str): The JSON path expression to the requested value.
+    """
+
+    step_name: str = attr.ib()
+    property_file: Union[PropertyFile, str] = attr.ib()
+    json_path: str = attr.ib()
+
+    @property
+    def expr(self):
+        """The expression dict for a `JsonGet` function."""
+        if not isinstance(self.step_name, str):
+            raise ValueError("Please give step name as a string")
+
+        if isinstance(self.property_file, PropertyFile):
+            name = self.property_file.name
+        else:
+            name = self.property_file
+        return {
+            "Std:JsonGet": {
+                "PropertyFile": {"Get": f"Steps.{self.step_name}.PropertyFiles.{name}"},
+                "Path": self.json_path,
+            }
+        }

tests/unit/sagemaker/workflow/test_functions.py

@@ -14,13 +14,13 @@
 from __future__ import absolute_import
 
 from sagemaker.workflow.execution_variables import ExecutionVariables
-from sagemaker.workflow.functions import Join
+from sagemaker.workflow.functions import Join, JsonGet
 from sagemaker.workflow.parameters import (
     ParameterFloat,
     ParameterInteger,
     ParameterString,
 )
-from sagemaker.workflow.properties import Properties
+from sagemaker.workflow.properties import Properties, PropertyFile
 
 
 def test_join_primitives_default_on():
@@ -66,3 +66,34 @@ def test_join_expressions():
             ],
         },
     }
+
+
+def test_json_get_expressions():
+
+    assert JsonGet(
+        step_name="my-step",
+        property_file="my-property-file",
+        json_path="my-json-path",
+    ).expr == {
+        "Std:JsonGet": {
+            "PropertyFile": {"Get": "Steps.my-step.PropertyFiles.my-property-file"},
+            "Path": "my-json-path",
+        },
+    }
+
+    property_file = PropertyFile(
+        name="name",
+        output_name="result",
+        path="output",
+    )
+
+    assert JsonGet(
+        step_name="my-step",
+        property_file=property_file,
+        json_path="my-json-path",
+    ).expr == {
+        "Std:JsonGet": {
+            "PropertyFile": {"Get": "Steps.my-step.PropertyFiles.name"},
+            "Path": "my-json-path",
+        },
+    }