feat(sdk): compile JSON with formatting (#7712)

* clean files

* implement formatted write

* add tests

sdk/python/kfp/utils/ir_utils.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Any
-import warnings
 import json
+import warnings
+from typing import Any, Dict
+
 import yaml
 
 
@@ -29,20 +30,19 @@ def _write_ir_to_file(ir_dict: Dict[str, Any], output_file: str) -> None:
         ValueError: If output file path is not JSON or YAML.
     """
 
-    if output_file.endswith(".json"):
+    if output_file.endswith('.json'):
         warnings.warn(
-            ("Compiling to JSON is deprecated and will be "
-             "removed in a future version. Please compile to a YAML file by "
-             "providing a file path with a .yaml extension instead."),
+            ('Compiling to JSON is deprecated and will be '
+             'removed in a future version. Please compile to a YAML file by '
+             'providing a file path with a .yaml extension instead.'),
             category=DeprecationWarning,
             stacklevel=2,
         )
-        ir_json = json.dumps(ir_dict, sort_keys=True)
         with open(output_file, 'w') as json_file:
-            json_file.write(ir_json)
-    elif output_file.endswith((".yaml", ".yml")):
+            json.dump(ir_dict, json_file, indent=2, sort_keys=True)
+    elif output_file.endswith(('.yaml', '.yml')):
         with open(output_file, 'w') as yaml_file:
             yaml.dump(ir_dict, yaml_file, sort_keys=True)
     else:
         raise ValueError(
-            f'The output path {output_file} should end with ".yaml".')
+            f'The output path {output_file} should end with ".yaml".')

sdk/python/kfp/utils/ir_utils_test.py

@@ -26,16 +26,17 @@
 
 import os
 import tempfile
+import textwrap
 import unittest
 
 import yaml
 from kfp.utils import ir_utils
 
-json_dict = {"key": "val", "list": ["1", 2, 3.0]}
+json_dict = {'key': 'val', 'list': ['1', 2, 3.0]}
 
 
 def load_from_file(filepath: str) -> str:
-    with open(filepath, "r") as f:
+    with open(filepath) as f:
         return yaml.safe_load(f)
 
 
@@ -57,11 +58,26 @@ class TestWriteIrToFile(unittest.TestCase):
 
     def test_json(self):
         with tempfile.TemporaryDirectory() as tempdir, self.assertWarnsRegex(
-                DeprecationWarning, r"Compiling to JSON is deprecated"):
+                DeprecationWarning, r'Compiling to JSON is deprecated'):
             temp_filepath = os.path.join(tempdir, 'output.json')
             ir_utils._write_ir_to_file(json_dict, temp_filepath)
-            actual = load_from_file(temp_filepath)
-        self.assertEqual(actual, json_dict)
+
+            with open(temp_filepath) as f:
+                actual_contents = f.read()
+
+            actual_dict = load_from_file(temp_filepath)
+
+        expected_contents = textwrap.dedent("""\
+                        {
+                          "key": "val",
+                          "list": [
+                            "1",
+                            2,
+                            3.0
+                          ]
+                        }""")
+        self.assertEqual(actual_contents, expected_contents)
+        self.assertEqual(actual_dict, json_dict)
 
     def test_incorrect_extension(self):
         with tempfile.TemporaryDirectory() as tempdir, self.assertRaisesRegex(