github.com/abemedia/go-don@v0.2.2-0.20240329015135-be88e32bb73b/encoding/encode.go (about)

     1  package encoding
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  
     7  	"github.com/valyala/fasthttp"
     8  )
     9  
    10  type (
    11  	Marshaler        = func(v any) ([]byte, error)
    12  	ContextMarshaler = func(ctx context.Context, v any) ([]byte, error)
    13  	ResponseEncoder  = func(ctx *fasthttp.RequestCtx, v any) error
    14  )
    15  
    16  type EncoderConstraint interface {
    17  	Marshaler | ContextMarshaler | ResponseEncoder
    18  }
    19  
    20  // RegisterEncoder registers a response encoder on a given media type.
    21  func RegisterEncoder[T EncoderConstraint](enc T, mime string, aliases ...string) {
    22  	switch e := any(enc).(type) {
    23  	case Marshaler:
    24  		encoders[mime] = func(ctx *fasthttp.RequestCtx, v any) error {
    25  			b, err := e(v)
    26  			if err != nil {
    27  				return err
    28  			}
    29  			ctx.Response.SetBodyRaw(b)
    30  			return nil
    31  		}
    32  
    33  	case ContextMarshaler:
    34  		encoders[mime] = func(ctx *fasthttp.RequestCtx, v any) error {
    35  			b, err := e(ctx, v)
    36  			if err != nil {
    37  				return err
    38  			}
    39  			ctx.Response.SetBodyRaw(b)
    40  			return nil
    41  		}
    42  
    43  	case ResponseEncoder:
    44  		encoders[mime] = e
    45  	}
    46  
    47  	for _, alias := range aliases {
    48  		encoders[alias] = encoders[mime]
    49  	}
    50  }
    51  
    52  // GetEncoder returns the response encoder for a given media type.
    53  func GetEncoder(mime string) ResponseEncoder {
    54  	mimeParts := strings.Split(mime, ",")
    55  	for _, part := range mimeParts {
    56  		if enc, ok := encoders[part]; ok {
    57  			return enc
    58  		}
    59  	}
    60  	return nil
    61  }
    62  
    63  var encoders = map[string]ResponseEncoder{}