From 9c3c1fc0106ed2174f1a7c4a322a08f0a553bc1e Mon Sep 17 00:00:00 2001
From: Fabian Meumertzheim <fabian@meumertzhe.im>
Date: Fri, 17 Mar 2023 14:24:45 +0100
Subject: [PATCH] Let buildozer preserve type-specific formatting (#1148)

Before this commit, when applying buildozer to e.g. MODULE.bazel file,
it would still apply BUILD-specific formatting rules.
---
 buildozer/buildozer_test.sh | 31 +++++++++++++++++++++++++++++++
 edit/buildozer.go           |  2 +-
 edit/default_buildifier.go  |  2 +-
 edit/edit.go                |  9 ++++++---
 edit/safe/buildifier.go     |  2 +-
 5 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/buildozer/buildozer_test.sh b/buildozer/buildozer_test.sh
index f474e7357..51350c54b 100755
--- a/buildozer/buildozer_test.sh
+++ b/buildozer/buildozer_test.sh
@@ -2029,4 +2029,35 @@ EOF
 )' pkg2
 }
 
+function test_module_bazel() {
+  cat > MODULE.bazel <<EOF
+module(
+    name = "foo",
+    version = "0.27.0",
+)
+
+bazel_dep(name = "gazelle", version = "0.30.0")
+
+go_deps = use_extension("@gazelle//:extensions.bzl", "go_deps")
+go_deps.from_file(go_mod = "//:go.mod")
+use_repo(go_deps, "com_example_foo")
+EOF
+
+  cat > MODULE.bazel.expected <<EOF
+module(
+    name = "foo",
+    version = "0.27.0",
+)
+
+bazel_dep(name = "gazelle", version = "0.30.0")
+
+go_deps = use_extension("@gazelle//:extensions.bzl", "go_deps")
+go_deps.from_file(go_mod = "//:go.mod")
+EOF
+
+  $buildozer 'delete' //MODULE.bazel:%10
+  diff -u MODULE.bazel.expected MODULE.bazel || fail "Output didn't match"
+}
+
+
 run_suite "buildozer tests"
diff --git a/edit/buildozer.go b/edit/buildozer.go
index 690ab9964..ff9c615c7 100644
--- a/edit/buildozer.go
+++ b/edit/buildozer.go
@@ -992,7 +992,7 @@ func rewrite(opts *Options, commandsForFile commandsForFile) *rewriteResult {
 		}
 	}
 
-	f, err := build.ParseBuild(name, data)
+	f, err := build.Parse(name, data)
 	if err != nil {
 		return &rewriteResult{file: name, errs: []error{err}}
 	}
diff --git a/edit/default_buildifier.go b/edit/default_buildifier.go
index e7bf24743..7b015a57b 100644
--- a/edit/default_buildifier.go
+++ b/edit/default_buildifier.go
@@ -20,7 +20,7 @@ func (b *defaultBuildifier) Buildify(opts *Options, f *build.File) ([]byte, erro
 		// value is a chunk of code, like "f(x)". The AST should be printed and
 		// re-read to parse such expressions correctly.
 		contents := build.Format(f)
-		newF, err := build.ParseBuild(f.Path, []byte(contents))
+		newF, err := build.Parse(f.Path, contents)
 		if err != nil {
 			return nil, err
 		}
diff --git a/edit/edit.go b/edit/edit.go
index ee70e874c..c500b3e34 100644
--- a/edit/edit.go
+++ b/edit/edit.go
@@ -172,6 +172,9 @@ func PackageDeclaration(f *build.File) *build.Rule {
 // This might appear because of a buildozer transformation (e.g. when removing a package
 // attribute). Removing it is required for the file to be valid.
 func RemoveEmptyPackage(f *build.File) *build.File {
+	if f.Type != build.TypeBuild {
+		return f
+	}
 	var all []build.Expr
 	for _, stmt := range f.Stmt {
 		if isEmptyPackage(stmt) {
@@ -315,7 +318,7 @@ func DeleteRule(f *build.File, rule *build.Rule) *build.File {
 		}
 		all = append(all, stmt)
 	}
-	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: build.TypeBuild}
+	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: f.Type}
 }
 
 // DeleteRuleByName returns the AST without the rules that have the
@@ -333,7 +336,7 @@ func DeleteRuleByName(f *build.File, name string) *build.File {
 			all = append(all, stmt)
 		}
 	}
-	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: build.TypeBuild}
+	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: f.Type}
 }
 
 // DeleteRuleByKind removes the rules of the specified kind from the AST.
@@ -351,7 +354,7 @@ func DeleteRuleByKind(f *build.File, kind string) *build.File {
 			all = append(all, stmt)
 		}
 	}
-	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: build.TypeBuild}
+	return &build.File{Path: f.Path, Comments: f.Comments, Stmt: all, Type: f.Type}
 }
 
 // AllLists returns all the lists concatenated in an expression.
diff --git a/edit/safe/buildifier.go b/edit/safe/buildifier.go
index ffaa87a6b..b39a1c482 100644
--- a/edit/safe/buildifier.go
+++ b/edit/safe/buildifier.go
@@ -24,7 +24,7 @@ func (b *buildifier) Buildify(_ *edit.Options, f *build.File) ([]byte, error) {
 	// value is a chunk of code, like "f(x)". The AST should be printed and
 	// re-read to parse such expressions correctly.
 	contents := build.Format(f)
-	newF, err := build.ParseBuild(f.Path, []byte(contents))
+	newF, err := build.Parse(f.Path, contents)
 	if err != nil {
 		return nil, err
 	}