create helper for sorting python env; clean up diff output

Author: themisvaltinos

sqlmesh/core/audit/definition.py

@@ -16,7 +16,8 @@
     default_catalog_validator,
     depends_on_validator,
     expression_validator,
-    python_env_payloads,
+    sort_python_env,
+    sorted_python_env_payloads,
 )
 from sqlmesh.core.model.common import make_python_env, single_value_or_tuple
 from sqlmesh.core.node import _Node
@@ -240,7 +241,7 @@ def depends_on(self) -> t.Set[str]:
     @property
     def sorted_python_env(self) -> t.List[t.Tuple[str, Executable]]:
         """Returns the python env sorted by executable kind and then var name."""
-        return sorted(self.python_env.items(), key=lambda x: (x[1].kind, x[0]))
+        return sort_python_env(self.python_env)
 
     @property
     def data_hash(self) -> str:
@@ -341,7 +342,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(expressions=python_env_payloads(self.sorted_python_env))
+            python_env = d.PythonCode(expressions=sorted_python_env_payloads(self.python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

sqlmesh/core/context_diff.py

@@ -19,7 +19,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import get_console
 from sqlmesh.core.macros import RuntimeStage
-from sqlmesh.core.model.common import python_env_payloads
+from sqlmesh.core.model.common import sorted_python_env_payloads
 from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotTableInfo
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
@@ -337,27 +337,22 @@ def extract_statements(statements: t.List[EnvironmentStatements], attr: str) ->
                 expr
                 for statement in statements
                 for expr in (
-                    python_env_payloads(
-                        sorted(statement.python_env.items(), key=lambda x: (x[1].kind, x[0]))
-                    )
+                    sorted_python_env_payloads(statement.python_env)
                     if attr == "python_env"
                     else getattr(statement, attr)
                 )
             ]
 
-        def format_diff(attribute: str) -> t.Optional[t.Tuple[str, str]]:
+        def compute_diff(attribute: str) -> t.Optional[t.Tuple[str, str]]:
             previous = extract_statements(self.previous_environment_statements, attribute)
             current = extract_statements(self.environment_statements, attribute)
 
             if previous == current:
                 return None
 
             diff_lines = list(ndiff(previous, current))
-            diff_text = (
-                f"=== {attribute} ===\n"
-                if not attribute == "python_env"
-                else "=== dependencies ===\n"
-            )
+            diff_text = attribute if not attribute == "python_env" else "dependencies"
+            diff_text += ":\n"
 
             if any(line.startswith(("-", "+")) for line in diff_lines):
                 diff_text += "  " + "\n  ".join(diff_lines) + "\n"
@@ -371,7 +366,7 @@ def format_diff(attribute: str) -> t.Optional[t.Tuple[str, str]]:
                 RuntimeStage.AFTER_ALL.value,
                 "python_env",
             ]
-            if (diff := format_diff(attribute)) is not None
+            if (diff := compute_diff(attribute)) is not None
         ]
 
     @property

sqlmesh/core/model/common.py

@@ -352,11 +352,16 @@ def depends_on(cls: t.Type, v: t.Any, info: ValidationInfo) -> t.Optional[t.Set[
     return v
 
 
-def python_env_payloads(sorted_python_env: t.List[t.Tuple[str, Executable]]) -> t.List[str]:
+def sort_python_env(python_env: t.Dict[str, Executable]) -> t.List[t.Tuple[str, Executable]]:
+    """Returns the python env sorted."""
+    return sorted(python_env.items(), key=lambda x: (x[1].kind, x[0]))
+
+
+def sorted_python_env_payloads(python_env: t.Dict[str, Executable]) -> t.List[str]:
     """Returns the payloads of the sorted python env."""
     return [
         v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-        for k, v in sorted_python_env
+        for k, v in sort_python_env(python_env)
     ]
 
 

sqlmesh/core/model/definition.py

@@ -29,8 +29,8 @@
     expression_validator,
     make_python_env,
     parse_dependencies,
-    python_env_payloads,
     single_value_or_tuple,
+    sorted_python_env_payloads,
 )
 from sqlmesh.core.model.meta import ModelMeta, FunctionCall
 from sqlmesh.core.model.kind import (
@@ -258,7 +258,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(expressions=python_env_payloads(self.sorted_python_env))
+            python_env = d.PythonCode(expressions=sorted_python_env_payloads(self.python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

tests/core/test_plan.py

@@ -2926,12 +2926,12 @@ def test_plan_environment_statements_diff(make_snapshot):
     output = console_output.getvalue()
     stripped = strip_ansi_codes(output)
     expected_output = (
-        "=== before_all ===\n"
+        "before_all:\n"
         "  + CREATE OR REPLACE TABLE table_1 AS SELECT 1\n"
         "  + @test_macro()\n\n"
-        "=== after_all ===\n"
+        "after_all:\n"
         "  + CREATE OR REPLACE TABLE table_2 AS SELECT 2\n\n"
-        "=== dependencies ===\n"
+        "dependencies:\n"
         "  + def test_macro(evaluator):\n"
         "    return 'one'"
     )