temp

[ci skip]

Author: benfdking

sqlmesh/core/linter/rules/builtin.py

@@ -215,7 +215,7 @@ def _standard_error_message(
             "Please register them in the external models file. This can be done by running 'sqlmesh create_external_models'.",
         )
 
-    def create_fix(self, model_name: str) -> t.Optional[TextEdit]:
+    def create_fix(self, model_name: str) -> t.Optional[Fix]:
         """
         Add an external model to the external models file.
         - If no external models file exists, it will create one with the model.
@@ -227,30 +227,40 @@ def create_fix(self, model_name: str) -> t.Optional[TextEdit]:
 
         external_models_path = root / EXTERNAL_MODELS_YAML
         if not external_models_path.exists():
-            return TextEdit(
-                path=external_models_path,
-                range=Range(start=Position(0, 0), end=Position(0, 0)),
-                new_text=f"- name: {model_name}\n",
+            return Fix(
+                title="Add external model",
+                edits=[
+                    TextEdit(
+                        path=external_models_path,
+                        range=Range(start=Position(0, 0), end=Position(0, 0)),
+                        new_text=f"- name: {model_name}\n",
+                    )
+                ],
             )
 
         # Figure out the position to insert the new external model at the end of the file, whether
         # needs new line or not.
         with open(external_models_path, "r", encoding="utf-8") as file:
-            lines = file.readlines()
+            lines = file.read()
 
-        if lines and lines[-1].strip():
-            # If the last line is not empty, we need to add a new line before adding the new model.
-            position = Position(line=len(lines), character=0)
-            new_text = f'\n- name: "{model_name}"\n'
+        # If file ends in newline, we can add the new model directly.
+        if not lines.endswith("\n"):
+            new_text = f"\n- name: {model_name}\n"
         else:
-            # If the last line is empty, we can add the new model directly.
-            position = Position(line=len(lines) - 1, character=0)
-            new_text = f'- name: "{model_name}"\n'
-
-        return TextEdit(
-            path=external_models_path,
-            range=Range(start=position, end=position),
-            new_text=new_text,
+            new_text = f"- name: {model_name}\n"
+
+        split_lines = lines.splitlines()
+        position = Position(line=len(split_lines) - 1, character=len(split_lines[-1]))
+
+        return Fix(
+            title="Add external model",
+            edits=[
+                TextEdit(
+                    path=external_models_path,
+                    range=Range(start=position, end=position),
+                    new_text=new_text,
+                )
+            ],
         )
 
 

tests/core/linter/test_builtin.py

@@ -53,22 +53,27 @@ def test_no_missing_external_models(tmp_path, copy_to_temp_path) -> None:
     )
     assert len(lint.fixes) == 1
     fix = lint.fixes[0]
-    fix_path = sushi_path / "external_models.yaml"
-    assert fix.path == fix_path
-    assert fix.new_text == '- name: "memory"."raw"."demographics"\n'
-    assert fix.range == Range(
+    assert len(fix.edits) == 1
+    edit = fix.edits[0]
+    assert edit.new_text == '- name: "memory"."raw"."demographics"\n'
+    assert edit.range == Range(
         start=Position(line=0, character=0),
         end=Position(line=0, character=0),
     )
+    fix_path = sushi_path / "external_models.yaml"
+    assert edit.path == fix_path
 
 
-def test_no_missing_external_models_with_existing_file_ending_in_newline(tmp_path, copy_to_temp_path) -> None:
+def test_no_missing_external_models_with_existing_file_ending_in_newline(
+    tmp_path, copy_to_temp_path
+) -> None:
     sushi_paths = copy_to_temp_path("examples/sushi")
     sushi_path = sushi_paths[0]
 
     # Overwrite the external_models.yaml file to end with a random file and a newline
+    os.remove(sushi_path / "external_models.yaml")
     with open(sushi_path / "external_models.yaml", "w") as f:
-        f.write("- name: \"memory\".\"raw\".\"test\"\n")
+        f.write("- name: memory.raw.test\n")
 
     # Override the config.py to turn on lint
     with open(sushi_path / "config.py", "r") as f:
@@ -105,21 +110,27 @@ def test_no_missing_external_models_with_existing_file_ending_in_newline(tmp_pat
     )
     assert len(lint.fixes) == 1
     fix = lint.fixes[0]
-    fix_path = sushi_path / "external_models.yaml"
-    assert fix.path == fix_path
-    assert fix.new_text == '- name: "memory"."raw"."demographics"\n'
-    assert fix.range == Range(
-        start=Position(line=0, character=0),
-        end=Position(line=0, character=0),
+    assert len(fix.edits) == 1
+    edit = fix.edits[0]
+    assert edit.new_text == '- name: "memory"."raw"."demographics"\n'
+    assert edit.range == Range(
+        start=Position(line=0, character=23),
+        end=Position(line=0, character=23),
     )
+    fix_path = sushi_path / "external_models.yaml"
+    assert edit.path == fix_path
 
-def test_no_missing_external_models_with_existing_file_not_ending_in_newline(tmp_path, copy_to_temp_path) -> None:
+
+def test_no_missing_external_models_with_existing_file_not_ending_in_newline(
+    tmp_path, copy_to_temp_path
+) -> None:
     sushi_paths = copy_to_temp_path("examples/sushi")
     sushi_path = sushi_paths[0]
 
     # Overwrite the external_models.yaml file to end with a random file and a newline
+    os.remove(sushi_path / "external_models.yaml")
     with open(sushi_path / "external_models.yaml", "w") as f:
-        f.write("- name: \"memory\".\"raw\".\"test\"")
+        f.write("- name: memory.raw.test")
 
     # Override the config.py to turn on lint
     with open(sushi_path / "config.py", "r") as f:
@@ -156,10 +167,12 @@ def test_no_missing_external_models_with_existing_file_not_ending_in_newline(tmp
     )
     assert len(lint.fixes) == 1
     fix = lint.fixes[0]
-    fix_path = sushi_path / "external_models.yaml"
-    assert fix.path == fix_path
-    assert fix.new_text == '- name: "memory"."raw"."demographics"\n'
-    assert fix.range == Range(
-        start=Position(line=0, character=0),
-        end=Position(line=0, character=0),
+    assert len(fix.edits) == 1
+    edit = fix.edits[0]
+    assert edit.new_text == '\n- name: "memory"."raw"."demographics"\n'
+    assert edit.range == Range(
+        start=Position(line=0, character=23),
+        end=Position(line=0, character=23),
     )
+    fix_path = sushi_path / "external_models.yaml"
+    assert edit.path == fix_path