Refactors

Author: themisvaltinos

sqlmesh/core/audit/definition.py

@@ -16,6 +16,7 @@
     default_catalog_validator,
     depends_on_validator,
     expression_validator,
+    python_env_payloads,
 )
 from sqlmesh.core.model.common import make_python_env, single_value_or_tuple
 from sqlmesh.core.node import _Node
@@ -340,12 +341,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(
-                expressions=[
-                    v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-                    for k, v in self.sorted_python_env
-                ]
-            )
+            python_env = d.PythonCode(expressions=python_env_payloads(self.sorted_python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

sqlmesh/core/console.py

@@ -1312,8 +1312,8 @@ def show_model_difference_summary(
 
         if context_diff.has_environment_statements_changes:
             self._print("[bold]Environment statements:\n")
-            for changes in context_diff.environment_statements_diff():
-                self._print(changes)
+            for type, diff in context_diff.environment_statements_diff():
+                self._print(Syntax(diff, type, line_numbers=False))
 
         self._show_summary_tree_for(
             context_diff,
@@ -2464,8 +2464,8 @@ def show_model_difference_summary(
 
         if context_diff.has_environment_statements_changes:
             self._print("[bold]Environment statements:\n")
-            for changes in context_diff.environment_statements_diff():
-                self._print(changes)
+            for _, diff in context_diff.environment_statements_diff():
+                self._print(diff)
 
         added_snapshots = {context_diff.snapshots[s_id] for s_id in context_diff.added}
         added_snapshot_models = {s for s in added_snapshots if s.is_model}
@@ -2979,8 +2979,8 @@ def show_model_difference_summary(
 
         if context_diff.has_environment_statements_changes:
             self._write("Environment statements:\n")
-            for changes in context_diff.environment_statements_diff():
-                self._write(changes)
+            for _, diff in context_diff.environment_statements_diff():
+                self._write(diff)
 
         for added in context_diff.new_snapshots:
             self._write(f"  Added: {added}")

sqlmesh/core/context_diff.py

@@ -16,10 +16,10 @@
 import typing as t
 from difflib import ndiff
 from functools import cached_property
-from rich.syntax import Syntax
 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.snapshot import Snapshot, SnapshotId, SnapshotTableInfo
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
@@ -331,58 +331,47 @@ def requirements_diff(self) -> str:
             )
         )
 
-    def environment_statements_diff(self) -> t.List[Syntax]:
-        PYTHON_ENV = "python_env"
-
-        def extract_python_env(python_env: t.Dict[str, Executable]) -> t.List[str]:
-            return [
-                v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-                for k, v in python_env.items()
-            ]
-
+    def environment_statements_diff(self) -> t.List[t.Tuple[str, str]]:
         def extract_statements(statements: t.List[EnvironmentStatements], attr: str) -> t.List[str]:
             return [
                 expr
                 for statement in statements
                 for expr in (
-                    extract_python_env(statement.python_env)
-                    if attr == PYTHON_ENV
+                    python_env_payloads(
+                        sorted(statement.python_env.items(), key=lambda x: (x[1].kind, x[0]))
+                    )
+                    if attr == "python_env"
                     else getattr(statement, attr)
                 )
             ]
 
-        def format_diff(attribute: str) -> t.Optional[Syntax]:
+        def format_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
+                if not attribute == "python_env"
                 else "=== dependencies ===\n"
             )
 
-            diff_lines = list(ndiff(previous, current))
             if any(line.startswith(("-", "+")) for line in diff_lines):
                 diff_text += "  " + "\n  ".join(diff_lines) + "\n"
 
-            return Syntax(
-                diff_text, "python" if attribute == PYTHON_ENV else "sql", line_numbers=False
-            )
+            return "python" if attribute == "python_env" else "sql", diff_text
 
         return [
             diff
-            for diff in (
-                format_diff(attribute)
-                for attribute in [
-                    RuntimeStage.BEFORE_ALL.value,
-                    RuntimeStage.AFTER_ALL.value,
-                    PYTHON_ENV,
-                ]
-            )
-            if diff is not None
+            for attribute in [
+                RuntimeStage.BEFORE_ALL.value,
+                RuntimeStage.AFTER_ALL.value,
+                "python_env",
+            ]
+            if (diff := format_diff(attribute)) is not None
         ]
 
     @property

sqlmesh/core/model/common.py

@@ -352,6 +352,14 @@ 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]:
+    """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
+    ]
+
+
 expression_validator: t.Callable = field_validator(
     "query",
     "expressions_",

sqlmesh/core/model/definition.py

@@ -29,6 +29,7 @@
     expression_validator,
     make_python_env,
     parse_dependencies,
+    python_env_payloads,
     single_value_or_tuple,
 )
 from sqlmesh.core.model.meta import ModelMeta, FunctionCall
@@ -257,12 +258,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(
-                expressions=[
-                    v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-                    for k, v in self.sorted_python_env
-                ]
-            )
+            python_env = d.PythonCode(expressions=python_env_payloads(self.sorted_python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

tests/core/test_plan.py

@@ -2921,22 +2921,18 @@ def test_plan_environment_statements_diff(make_snapshot):
     assert context_diff.has_environment_statements_changes
 
     console_output, terminal_console = create_test_console()
-
-    for stmt in context_diff.environment_statements_diff():
-        terminal_console._print(stmt)
+    for _, diff in context_diff.environment_statements_diff():
+        terminal_console._print(diff)
     output = console_output.getvalue()
     stripped = strip_ansi_codes(output)
-
     expected_output = (
-        "=== before_all ===                                                              \n"
-        "  + CREATE OR REPLACE TABLE table_1 AS SELECT 1                                 \n"
-        "  + @test_macro()                                                               \n"
-        "                                                                                \n"
-        "=== after_all ===                                                               \n"
-        "  + CREATE OR REPLACE TABLE table_2 AS SELECT 2                                 \n"
-        "                                                                                \n"
-        "=== dependencies ===                                                            \n"
-        "  + def test_macro(evaluator):                                                  \n"
+        "=== before_all ===\n"
+        "  + CREATE OR REPLACE TABLE table_1 AS SELECT 1\n"
+        "  + @test_macro()\n\n"
+        "=== after_all ===\n"
+        "  + CREATE OR REPLACE TABLE table_2 AS SELECT 2\n\n"
+        "=== dependencies ===\n"
+        "  + def test_macro(evaluator):\n"
         "    return 'one'"
     )
     assert stripped == expected_output