Expose gopatch API as patch package

e2e_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/rogpeppe/go-internal/txtar"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/uber-go/gopatch/patch"
 	"go.uber.org/multierr"
 )
 
@@ -159,6 +160,19 @@ func runIntegrationTest(t *testing.T, testFile string) {
 				assert.Equal(t, string(tt.WantComment), stderr)
 			})
 
+			t.Run("api", func(t *testing.T) {
+				if stdin == nil {
+					t.Skipf("this test case is not single patch: %v", args)
+				}
+				got, err := os.ReadFile(filepath.Join(dir, filePath))
+				require.NoError(t, err, "failed to read %q", filePath)
+				patchFile, err := patch.Parse(filePath, stdin)
+				require.NoError(t, err, "failed to parse patch file %q", filePath)
+				actual, err := patchFile.Apply(filePath, got)
+				require.NoError(t, err, "failed to apply patch file %q", filePath)
+				assert.Equal(t, string(tt.Want), string(actual))
+			})
+
 			_, _, err := run(filePath)
 			require.NoError(t, err, "could not run patch")
 

patch/gopatch.go

@@ -0,0 +1,133 @@
+package patch
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"golang.org/x/tools/imports"
+	"sort"
+
+	"github.com/uber-go/gopatch/internal/astdiff"
+	"github.com/uber-go/gopatch/internal/engine"
+	"github.com/uber-go/gopatch/internal/parse"
+)
+
+// File is a patch difference file that can be applied to Go file.
+type File struct {
+	fset *token.FileSet
+	prog *engine.Program
+}
+
+// Parse the patch file and creates data that can be applied to the Go file.
+func Parse(patchFileName string, src []byte) (*File, error) {
+	fset := token.NewFileSet()
+
+	astProg, err := parse.Parse(fset, patchFileName, src)
+	if err != nil {
+		return nil, fmt.Errorf("parse: %w", err)
+	}
+	prog, err := engine.Compile(fset, astProg)
+	if err != nil {
+		return nil, fmt.Errorf("compile: %w", err)
+	}
+
+	return &File{fset: fset, prog: prog}, nil
+}
+
+// Apply takes the Go file name and its contents and returns a Go file with the patch applied.
+func (f *File) Apply(filename string, src []byte) ([]byte, error) {
+	base, err := parser.ParseFile(f.fset, filename, src, parser.AllErrors|parser.ParseComments)
+	if err != nil {
+		return nil, fmt.Errorf("could not parse %q: %w", filename, err)
+	}
+
+	snap := astdiff.Before(base, ast.NewCommentMap(f.fset, base, base.Comments))
+
+	var fout *ast.File
+	var retErr error
+	for _, c := range f.prog.Changes {
+		d, ok := c.Match(base)
+		if !ok {
+			// This patch didn't modify the file. Try the next one.
+			continue
+		}
+
+		cl := engine.NewChangelog()
+
+		fout, err = c.Replace(d, cl)
+		if err != nil {
+			retErr = errors.Join(retErr, err)
+			continue
+		}
+
+		snap = snap.Diff(fout, cl)
+		cleanupFilePos(f.fset.File(fout.Pos()), cl, fout.Comments)
+	}
+
+	if retErr != nil {
+		return nil, retErr
+	}
+
+	if fout == nil {
+		return src, nil
+	}
+
+	var out bytes.Buffer
+	err = format.Node(&out, f.fset, fout)
+	if err != nil {
+		return nil, err
+	}
+
+	bs := out.Bytes()
+	bs, err = imports.Process(filename, bs, &imports.Options{
+		Comments:   true,
+		TabIndent:  true,
+		TabWidth:   8,
+		FormatOnly: true,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return bs, nil
+}
+
+func cleanupFilePos(tfile *token.File, cl engine.Changelog, comments []*ast.CommentGroup) {
+	linesToDelete := make(map[int]struct{})
+	for _, dr := range cl.ChangedIntervals() {
+		if dr.Start == token.NoPos {
+			continue
+		}
+
+		for i := tfile.Line(dr.Start); i < tfile.Line(dr.End); i++ {
+			if i > 0 {
+				linesToDelete[i] = struct{}{}
+			}
+		}
+
+		// Remove comments in the changed sections of the code.
+		for _, cg := range comments {
+			var list []*ast.Comment
+			for _, c := range cg.List {
+				if c.Pos() >= dr.Start && c.End() <= dr.End {
+					continue
+				}
+				list = append(list, c)
+			}
+			cg.List = list
+		}
+	}
+
+	lines := make([]int, 0, len(linesToDelete))
+	for i := range linesToDelete {
+		lines = append(lines, i)
+	}
+	sort.Ints(lines)
+	for i := len(lines) - 1; i >= 0; i-- {
+		tfile.MergeLine(lines[i])
+	}
+}