refactor: put structure in place to support CodeAction behaviour

.version

@@ -1 +1 @@
-0.2.731
+0.2.732

cmd/templ/lspcmd/lsp_test.go

@@ -319,6 +319,129 @@ func TestHover(t *testing.T) {
 	}
 }
 
+func TestCodeAction(t *testing.T) {
+	if testing.Short() {
+		return
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	log, _ := zap.NewProduction()
+
+	ctx, appDir, _, server, teardown, err := Setup(ctx, log)
+	if err != nil {
+		t.Fatalf("failed to setup test: %v", err)
+	}
+	defer teardown(t)
+	defer cancel()
+
+	templFile, err := os.ReadFile(appDir + "/templates.templ")
+	if err != nil {
+		t.Fatalf("failed to read file %q: %v", appDir+"/templates.templ", err)
+	}
+	err = server.DidOpen(ctx, &protocol.DidOpenTextDocumentParams{
+		TextDocument: protocol.TextDocumentItem{
+			URI:        uri.URI("file://" + appDir + "/templates.templ"),
+			LanguageID: "templ",
+			Version:    1,
+			Text:       string(templFile),
+		},
+	})
+	if err != nil {
+		t.Errorf("failed to register open file: %v", err)
+		return
+	}
+	log.Info("Calling codeAction")
+
+	tests := []struct {
+		line        int
+		replacement string
+		cursor      string
+		assert      func(t *testing.T, hr []protocol.CodeAction) (msg string, ok bool)
+	}{
+		{
+			line:        25,
+			replacement: `var s = Struct{}`,
+			cursor:      `              ^`,
+			assert: func(t *testing.T, actual []protocol.CodeAction) (msg string, ok bool) {
+				var expected []protocol.CodeAction
+				// To support code actions, update cmd/templ/lspcmd/proxy/server.go and add the
+				// Title (e.g. Organize Imports, or Fill Struct) to the supportedCodeActions map.
+
+				// Some Code Actions are simple edits, so all that is needed is for the server
+				// to remap the source code positions.
+
+				// However, other Code Actions are commands, where the arguments must be rewritten
+				// and will need to be handled individually.
+				if diff := lspdiff.CodeAction(expected, actual); diff != "" {
+					return fmt.Sprintf("unexpected codeAction: %v", diff), false
+				}
+				return "", true
+			},
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
+			// Put the file back to the initial point.
+			err = server.DidChange(ctx, &protocol.DidChangeTextDocumentParams{
+				TextDocument: protocol.VersionedTextDocumentIdentifier{
+					TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+						URI: uri.URI("file://" + appDir + "/templates.templ"),
+					},
+					Version: int32(i + 2),
+				},
+				ContentChanges: []protocol.TextDocumentContentChangeEvent{
+					{
+						Range: nil,
+						Text:  string(templFile),
+					},
+				},
+			})
+			if err != nil {
+				t.Errorf("failed to change file: %v", err)
+				return
+			}
+
+			// Give CI/CD pipeline executors some time because they're often quite slow.
+			var ok bool
+			var msg string
+			for i := 0; i < 3; i++ {
+				lspCharIndex, err := runeIndexToUTF8ByteIndex(test.replacement, len(test.cursor)-1)
+				if err != nil {
+					t.Error(err)
+				}
+				actual, err := server.CodeAction(ctx, &protocol.CodeActionParams{
+					TextDocument: protocol.TextDocumentIdentifier{
+						URI: uri.URI("file://" + appDir + "/templates.templ"),
+					},
+					Range: protocol.Range{
+						Start: protocol.Position{
+							Line:      uint32(test.line - 1),
+							Character: lspCharIndex,
+						},
+						End: protocol.Position{
+							Line:      uint32(test.line - 1),
+							Character: lspCharIndex + 1,
+						},
+					},
+				})
+				if err != nil {
+					t.Errorf("failed code action: %v", err)
+					return
+				}
+				msg, ok = test.assert(t, actual)
+				if !ok {
+					break
+				}
+				time.Sleep(time.Millisecond * 500)
+			}
+			if !ok {
+				t.Error(msg)
+			}
+		})
+	}
+}
+
 func runeIndexToUTF8ByteIndex(s string, runeIndex int) (lspChar uint32, err error) {
 	for i, r := range []rune(s) {
 		if i == runeIndex {

cmd/templ/lspcmd/lspdiff/lspdiff.go

@@ -15,6 +15,10 @@ func Hover(expected, actual protocol.Hover) string {
 	)
 }
 
+func CodeAction(expected, actual []protocol.CodeAction) string {
+	return cmp.Diff(expected, actual)
+}
+
 func CompletionList(expected, actual *protocol.CompletionList) string {
 	return cmp.Diff(expected, actual,
 		cmpopts.IgnoreFields(protocol.CompletionList{}, "IsIncomplete"),

cmd/templ/lspcmd/proxy/server.go

@@ -273,6 +273,8 @@ func (p *Server) SetTrace(ctx context.Context, params *lsp.SetTraceParams) (err
 	return p.Target.SetTrace(ctx, params)
 }
 
+var supportedCodeActions = map[string]bool{}
+
 func (p *Server) CodeAction(ctx context.Context, params *lsp.CodeActionParams) (result []lsp.CodeAction, err error) {
 	p.Log.Info("client -> server: CodeAction")
 	defer p.Log.Info("client -> server: CodeAction end")
@@ -281,13 +283,22 @@ func (p *Server) CodeAction(ctx context.Context, params *lsp.CodeActionParams) (
 		return p.Target.CodeAction(ctx, params)
 	}
 	templURI := params.TextDocument.URI
+	params.Range = p.convertTemplRangeToGoRange(templURI, params.Range)
 	params.TextDocument.URI = goURI
-	params.Range = p.convertTemplRangeToGoRange(params.TextDocument.URI, params.Range)
 	result, err = p.Target.CodeAction(ctx, params)
 	if err != nil {
 		return
 	}
+	var updatedResults []lsp.CodeAction
+	// Filter out commands that are not yet supported.
+	// For example, "Fill Struct" runs the `gopls.apply_fix` command.
+	// This command has a set of arguments, including Fix, Range and URI.
+	// However, these are just a map[string]any so for each command that we want to support,
+	// we need to know what the arguments are so that we can rewrite them.
 	for i := 0; i < len(result); i++ {
+		if !supportedCodeActions[result[i].Title] {
+			continue
+		}
 		r := result[i]
 		// Rewrite the Diagnostics range field.
 		for di := 0; di < len(r.Diagnostics); di++ {
@@ -301,11 +312,12 @@ func (p *Server) CodeAction(ctx context.Context, params *lsp.CodeActionParams) (
 					dc.Edits[ei].Range = p.convertGoRangeToTemplRange(templURI, dc.Edits[ei].Range)
 				}
 				dc.TextDocument.URI = templURI
+				r.Edit.DocumentChanges[dci] = dc
 			}
 		}
-		result[i] = r
+		updatedResults = append(updatedResults, r)
 	}
-	return
+	return updatedResults, nil
 }
 
 func (p *Server) CodeLens(ctx context.Context, params *lsp.CodeLensParams) (result []lsp.CodeLens, err error) {
@@ -555,6 +567,12 @@ func (p *Server) DidChange(ctx context.Context, params *lsp.DidChangeTextDocumen
 		return
 	}
 	w := new(strings.Builder)
+	// In future updates, we may pass `WithSkipCodeGeneratedComment` to the generator.
+	// This will enable a number of actions within gopls that it doesn't currently apply because
+	// it recognises templ code as being auto-generated.
+	//
+	// This change would increase the surface area of gopls that we use, so may surface a number of issues
+	// if enabled.
 	sm, _, err := generator.Generate(template, w)
 	if err != nil {
 		p.Log.Error("generate failure", zap.Error(err))

cmd/templ/testproject/testdata/templates.templ

@@ -17,3 +17,9 @@ templ Page(count int) {
 }
 
 var nihao = "你好"
+
+type Struct struct {
+	Count int
+}
+
+var s = Struct{}

generator/generator.go

@@ -57,6 +57,16 @@ func WithExtractStrings() GenerateOpt {
 	}
 }
 
+// WithSkipCodeGeneratedComment skips the code generated comment at the top of the file.
+// gopls disables edit related functionality for generated files, so the templ LSP may
+// wish to skip generation of this comment so that gopls provides expected results.
+func WithSkipCodeGeneratedComment() GenerateOpt {
+	return func(g *generator) error {
+		g.skipCodeGeneratedComment = true
+		return nil
+	}
+}
+
 // Generate generates Go code from the input template file to w, and returns a map of the location of Go expressions in the template
 // to the location of the generated Go code in the output.
 func Generate(template parser.TemplateFile, w io.Writer, opts ...GenerateOpt) (sm *parser.SourceMap, literals string, err error) {
@@ -89,6 +99,8 @@ type generator struct {
 	generatedDate string
 	// fileName to include in error messages if string expressions return an error.
 	fileName string
+	// skipCodeGeneratedComment skips the code generated comment at the top of the file.
+	skipCodeGeneratedComment bool
 }
 
 func (g *generator) generate() (err error) {
@@ -116,7 +128,15 @@ func (g *generator) generate() (err error) {
 	return err
 }
 
+// See https://pkg.go.dev/cmd/go#hdr-Generate_Go_files_by_processing_source
+// Automatically generated files have a comment in the header that instructs the LSP
+// to stop operating.
 func (g *generator) writeCodeGeneratedComment() (err error) {
+	if g.skipCodeGeneratedComment {
+		// Write an empty comment so that the file is the same shape.
+		_, err = g.w.Write("//\n\n")
+		return err
+	}
 	_, err = g.w.Write("// Code generated by templ - DO NOT EDIT.\n\n")
 	return err
 }