github.com/TrueBlocks/trueblocks-core/src/apps/chifra@v0.0.0-20241022031540-b362680128f7/pkg/output/helpers/helpers.go (about)

     1  package outputHelpers
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  
     7  	"github.com/TrueBlocks/trueblocks-core/src/apps/chifra/internal/globals"
     8  	"github.com/TrueBlocks/trueblocks-core/src/apps/chifra/pkg/output"
     9  	"github.com/TrueBlocks/trueblocks-core/src/apps/chifra/pkg/rpc"
    10  	"github.com/TrueBlocks/trueblocks-core/src/apps/chifra/pkg/types"
    11  	"github.com/spf13/cobra"
    12  )
    13  
    14  var discardWriter = io.Writer(nil)
    15  
    16  func init() {
    17  	if os.Getenv("TB_DISCARD_WRITER") == "true" {
    18  		discardWriter = io.Discard
    19  	}
    20  }
    21  
    22  func PreRunWithJsonWriter(cmdName string, getOptions func() *globals.GlobalOptions) func(cmd *cobra.Command, args []string) {
    23  	return func(cmd *cobra.Command, args []string) {
    24  		opts := getOptions()
    25  		var outputWriter io.Writer
    26  		// Prepare the default writer (stdout or --output file)
    27  		outputWriter = os.Stdout
    28  		if opts.OutputFn != "" {
    29  			outputWriter = opts.GetOutputFileWriter()
    30  		} else if discardWriter != nil {
    31  			outputWriter = discardWriter
    32  		}
    33  
    34  		// If we need to output JSON, init JsonWriter...
    35  		if opts.Format == "json" {
    36  			jw := output.NewDefaultJsonWriter(outputWriter, true)
    37  			opts.Writer = jw
    38  		} else {
    39  			// ...or set the default writer as global writer for the current command
    40  			// invocation
    41  			opts.Writer = outputWriter
    42  		}
    43  	}
    44  }
    45  
    46  func PostRunWithJsonWriter(getOptions func() *globals.GlobalOptions) func(cmd *cobra.Command, args []string) {
    47  	return func(cmd *cobra.Command, args []string) {
    48  		opts := getOptions()
    49  		w := opts.Writer
    50  		// Try to cast the global writer to JsonWriter
    51  		jw, ok := w.(*output.JsonWriter)
    52  		if !ok {
    53  			// If it's not JsonWriter, do nothing
    54  			return
    55  		}
    56  		// If it is JsonWriter, close it (print closing brackets)
    57  		jw.Close()
    58  	}
    59  }
    60  
    61  // SetWriterForCommand sets the writer for currently running command, but only if
    62  // we are running with --file
    63  func SetWriterForCommand(cmdName string, opts *globals.GlobalOptions) {
    64  	// Try to cast the default writer to JsonWriter
    65  	jw, ok := opts.Writer.(*output.JsonWriter)
    66  	wantsJson := (opts.Format == "json")
    67  
    68  	// Global writer is set to JsonWriter, but this command wants to output
    69  	// a different format. We have to close JsonWriter so that closing brackets
    70  	// are printed and switch global writer to io.Writer
    71  	if ok && !wantsJson {
    72  		jw.Close()
    73  		// Set global writer to default
    74  		opts.Writer = os.Stdout
    75  		opts.Writer = opts.GetOutputFileWriter()
    76  		return
    77  	}
    78  
    79  	// Global writer is NOT JsonWriter, but this command wants to output
    80  	// JSON, so we have to create new JsonWriter and set is as the global
    81  	// writer
    82  	if !ok && wantsJson {
    83  		w := opts.GetOutputFileWriter()
    84  		jw := output.NewDefaultJsonWriter(w, true)
    85  		opts.Writer = jw
    86  	}
    87  
    88  	// Global writer is not JsonWriter and the command doesn't want to
    89  	// output JSON, so we only have to set global writer to a default one
    90  	if !ok && !wantsJson {
    91  		opts.Writer = opts.GetOutputFileWriter()
    92  	}
    93  }
    94  
    95  // InitJsonWriterApi inits JsonWriter for API responses
    96  func InitJsonWriterApi(cmdName string, w io.Writer, opts *globals.GlobalOptions) {
    97  	_, ok := opts.Writer.(*output.JsonWriter)
    98  	if opts.Format == "json" && !ok {
    99  		jw := output.NewDefaultJsonWriter(w, false)
   100  		jw.ShouldWriteMeta = true
   101  		jw.GetMeta = func() (*types.MetaData, error) {
   102  			chain := opts.Chain
   103  			conn := rpc.TempConnection(chain)
   104  			return conn.GetMetaData(opts.OutputOptions.TestMode)
   105  		}
   106  		opts.Writer = jw
   107  	}
   108  }
   109  
   110  // CloseJsonWriterIfNeededApi will close JsonWriter if the format is json
   111  func CloseJsonWriterIfNeededApi(cmdName string, err error, opts *globals.GlobalOptions) {
   112  	if opts.Format == "json" && err == nil {
   113  		opts.Writer.(*output.JsonWriter).Close()
   114  	}
   115  }