chore: Add Options to GenTreeDocs and move buf to io.Writer (#68)

Author: fmenezes

cobra2snooty.go

@@ -35,12 +35,12 @@ const (
 )
 
 // GenTreeDocs generates the docs for the full tree of commands.
-func GenTreeDocs(cmd *cobra.Command, dir string) error {
+func GenTreeDocs(cmd *cobra.Command, dir string, genDocOptions ...GenDocsOption) error {
 	for _, c := range cmd.Commands() {
 		if !c.IsAvailableCommand() || c.IsAdditionalHelpTopicCommand() {
 			continue
 		}
-		if err := GenTreeDocs(c, dir); err != nil {
+		if err := GenTreeDocs(c, dir, genDocOptions...); err != nil {
 			return err
 		}
 	}
@@ -53,7 +53,7 @@ func GenTreeDocs(cmd *cobra.Command, dir string) error {
 	}
 	defer f.Close()
 
-	return GenDocs(cmd, f)
+	return GenDocs(cmd, f, genDocOptions...)
 }
 
 const toc = `
@@ -198,7 +198,7 @@ func (s byName) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
 func (s byName) Less(i, j int) bool { return s[i].Name() < s[j].Name() }
 
 type GenDocsOptions struct {
-	exampleFormatter func(buf *bytes.Buffer, cmd *cobra.Command)
+	exampleFormatter ExampleFormatter
 	timeGetter       func() time.Time
 }
 

cobra2snooty_test.go

@@ -17,6 +17,7 @@ package cobra2snooty
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -317,9 +318,9 @@ func TestGenDocsSnapshots(t *testing.T) {
 				WithCustomTimeGetter(func() time.Time {
 					return time.Date(2025, 3, 5, 17, 0, 0, 0, time.UTC)
 				}),
-				WithCustomExampleFormatter(func(buf *bytes.Buffer, cmd *cobra.Command) {
-					_, _ = fmt.Fprintf(buf, "custom example for %s\n", cmd.Use)
-					_, _ = buf.WriteString(cmd.Example)
+				WithCustomExampleFormatter(func(w io.Writer, cmd *cobra.Command) {
+					_, _ = fmt.Fprintf(w, "custom example for %s\n", cmd.Use)
+					_, _ = w.Write([]byte(cmd.Example))
 				}),
 			},
 		},
@@ -334,10 +335,10 @@ func TestGenDocsSnapshots(t *testing.T) {
 				WithCustomTimeGetter(func() time.Time {
 					return time.Date(2025, 3, 5, 17, 0, 0, 0, time.UTC)
 				}),
-				WithCustomExampleFormatter(func(buf *bytes.Buffer, cmd *cobra.Command) {
-					_, _ = buf.WriteString("-- before example --\n")
-					DefaultExampleFormatter(buf, cmd)
-					_, _ = buf.WriteString("-- after example --\n")
+				WithCustomExampleFormatter(func(w io.Writer, cmd *cobra.Command) {
+					_, _ = w.Write([]byte("-- before example --\n"))
+					DefaultExampleFormatter(w, cmd)
+					_, _ = w.Write([]byte("-- after example --\n"))
 				}),
 			},
 		},

examples.go

@@ -15,8 +15,8 @@
 package cobra2snooty
 
 import (
-	"bytes"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -30,32 +30,34 @@ const (
 	identChar = " "
 )
 
-func DefaultExampleFormatter(buf *bytes.Buffer, cmd *cobra.Command) {
+type ExampleFormatter func(w io.Writer, cmd *cobra.Command)
+
+func DefaultExampleFormatter(w io.Writer, cmd *cobra.Command) {
 	if cmd.Example != "" {
-		printExamples(buf, cmd)
+		printExamples(w, cmd)
 	}
 }
 
-func WithCustomExampleFormatter(customFormatter func(buf *bytes.Buffer, cmd *cobra.Command)) func(options *GenDocsOptions) {
+func WithCustomExampleFormatter(customFormatter ExampleFormatter) func(options *GenDocsOptions) {
 	return func(options *GenDocsOptions) {
 		options.exampleFormatter = customFormatter
 	}
 }
 
-func printExamples(buf *bytes.Buffer, cmd *cobra.Command) {
+func printExamples(w io.Writer, cmd *cobra.Command) {
 	// Create example substrings
 	examplestrimmed := strings.TrimLeft(cmd.Example, " #")
 	examples := strings.Split(examplestrimmed, "# ")
-	buf.WriteString(examplesHeader)
+	_, _ = w.Write([]byte(examplesHeader))
 	// If it has an example, print the header, then print each in a code block.
 	for _, example := range examples[0:] {
 		comment := ""
 		if strings.Contains(cmd.Example, "#") {
 			comment = " #"
 		}
-		buf.WriteString(`.. code-block::
+		_, _ = w.Write([]byte(`.. code-block::
    :copyable: false
-`)
-		_, _ = fmt.Fprintf(buf, "\n  %s%s\n", comment, indentString(example, identChar))
+`))
+		_, _ = fmt.Fprintf(w, "\n  %s%s\n", comment, indentString(example, identChar))
 	}
 }