Feat: Improve the before all after all diff and include python env

Author: themisvaltinos

sqlmesh/core/console.py

@@ -1311,9 +1311,9 @@ def show_model_difference_summary(
             self._print(f"[bold]Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._print(
-                f"[bold]Environment statements:\n{context_diff.environment_statements_diff()}"
-            )
+            self._print("[bold]Environment statements:\n")
+            for changes in context_diff.environment_statements_diff():
+                self._print(changes)
 
         self._show_summary_tree_for(
             context_diff,
@@ -2463,7 +2463,9 @@ def show_model_difference_summary(
             self._print(f"Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._print(f"Environment statements:\n{context_diff.environment_statements_diff()}")
+            self._print("[bold]Environment statements:\n")
+            for changes in context_diff.environment_statements_diff():
+                self._print(changes)
 
         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}
@@ -2976,7 +2978,9 @@ def show_model_difference_summary(
             self._write(f"Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._write(f"Environment statements:\n{context_diff.environment_statements_diff()}")
+            self._write("Environment statements:\n")
+            for changes in context_diff.environment_statements_diff():
+                self._write(changes)
 
         for added in context_diff.new_snapshots:
             self._write(f"  Added: {added}")

sqlmesh/core/context_diff.py

@@ -16,6 +16,7 @@
 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
@@ -330,22 +331,55 @@ def requirements_diff(self) -> str:
             )
         )
 
-    def environment_statements_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 extract_statements(statements: t.List[EnvironmentStatements], attr: str) -> t.List[str]:
-            return [str(stmt) for statement in statements for stmt in getattr(statement, attr)]
-
-        def format_diff(runtime_stage: str) -> str:
-            previous = extract_statements(self.previous_environment_statements, runtime_stage)
-            current = extract_statements(self.environment_statements, runtime_stage)
-            return (
-                f"  {runtime_stage}:\n    " + "\n    ".join(ndiff(previous, current)) + "\n"
-                if previous or current
-                else ""
+            return [
+                expr
+                for statement in statements
+                for expr in (
+                    extract_python_env(statement.python_env)
+                    if attr == PYTHON_ENV
+                    else getattr(statement, attr)
+                )
+            ]
+
+        def format_diff(attribute: str) -> t.Optional[Syntax]:
+            previous = extract_statements(self.previous_environment_statements, attribute)
+            current = extract_statements(self.environment_statements, attribute)
+
+            if previous == current:
+                return None
+
+            diff_text = f"=== {attribute 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 format_diff(RuntimeStage.BEFORE_ALL.value) + format_diff(
-            RuntimeStage.AFTER_ALL.value
-        )
+        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
+        ]
 
     @property
     def environment_snapshots(self) -> t.List[SnapshotTableInfo]:

tests/core/test_plan.py

@@ -4,6 +4,8 @@
 from unittest.mock import patch
 
 import pytest
+from sqlmesh.utils.metaprogramming import Executable
+from tests.core.test_table_diff import create_test_console, strip_ansi_codes
 import time_machine
 from pytest_mock.plugin import MockerFixture
 from sqlglot import parse_one
@@ -2904,20 +2906,38 @@ def test_plan_environment_statements_diff(make_snapshot):
         previous_finalized_snapshots=None,
         environment_statements=[
             EnvironmentStatements(
-                before_all=["CREATE OR REPLACE TABLE table_1 AS SELECT 1"],
+                before_all=["CREATE OR REPLACE TABLE table_1 AS SELECT 1", "@test_macro()"],
                 after_all=["CREATE OR REPLACE TABLE table_2 AS SELECT 2"],
-                python_env={},
+                python_env={
+                    "test_macro": Executable(
+                        payload="def test_macro(evaluator):\n    return 'one'"
+                    ),
+                },
             )
         ],
     )
 
     assert context_diff.has_changes
     assert context_diff.has_environment_statements_changes
-    assert (
-        context_diff.environment_statements_diff()
-        == """  before_all:
-    + CREATE OR REPLACE TABLE table_1 AS SELECT 1
-  after_all:
-    + CREATE OR REPLACE TABLE table_2 AS SELECT 2
-"""
-    )
+
+    console_output, terminal_console = create_test_console()
+
+    for stmt in context_diff.environment_statements_diff():
+        terminal_console._print(stmt)
+    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"
+        "    return 'one'"
+    )
+    assert stripped == expected_output
+    console_output.close()