x/stream: ReadSource, ReadSourceLocal, ReadSourceFromURI

Author: xushiwei

stream/stream.go

@@ -17,16 +17,22 @@
 package stream
 
 import (
+	"bytes"
 	"errors"
 	"io"
 	"io/fs"
 	"os"
 	"strings"
+
+	"github.com/qiniu/x/byteutil"
 )
 
 var (
 	// ErrUnknownScheme is returned when an unknown scheme is encountered in a URL.
 	ErrUnknownScheme = errors.New("unknown scheme")
+
+	// ErrInvalidSource is returned when an invalid source is encountered.
+	ErrInvalidSource = errors.New("invalid source")
 )
 
 // -------------------------------------------------------------------------------------
@@ -43,28 +49,79 @@ func Register(scheme string, open OpenFunc) {
 	openers[scheme] = open
 }
 
-// Open opens a resource identified by the given URL.
+// Open opens a resource identified by the given URI.
 // It supports different schemes by utilizing registered open functions.
-// If the URL has no scheme, it is treated as a file path.
-func Open(url string) (io.ReadCloser, error) {
-	scheme := schemeOf(url)
+// If the URI has no scheme, it is treated as a file path.
+func Open(uri string) (io.ReadCloser, error) {
+	scheme := schemeOf(uri)
 	if scheme == "" {
-		return os.Open(url)
+		return os.Open(uri)
 	}
 	if open, ok := openers[scheme]; ok {
-		return open(url)
+		return open(uri)
 	}
-	return nil, &fs.PathError{Op: "stream.Open", Err: ErrUnknownScheme, Path: url}
+	return nil, &fs.PathError{Op: "stream.Open", Err: ErrUnknownScheme, Path: uri}
 }
 
-func schemeOf(url string) (scheme string) {
-	pos := strings.IndexAny(url, ":/")
+func schemeOf(uri string) (scheme string) {
+	pos := strings.IndexAny(uri, ":/")
 	if pos > 0 {
-		if url[pos] == ':' {
-			return url[:pos]
+		if uri[pos] == ':' {
+			return uri[:pos]
 		}
 	}
 	return ""
 }
 
 // -------------------------------------------------------------------------------------
+
+// ReadSource converts src to a []byte if possible; otherwise it returns an error.
+// Supported types for src are:
+//   - string (as content, NOT as filename)
+//   - []byte (as content)
+//   - *bytes.Buffer (as content)
+//   - io.Reader (as content)
+func ReadSource(src any) ([]byte, error) {
+	switch s := src.(type) {
+	case string:
+		return byteutil.Bytes(s), nil
+	case []byte:
+		return s, nil
+	case *bytes.Buffer:
+		// is io.Reader, but src is already available in []byte form
+		if s != nil {
+			return s.Bytes(), nil
+		}
+	case io.Reader:
+		return io.ReadAll(s)
+	}
+	return nil, ErrInvalidSource
+}
+
+// If src != nil, ReadSourceLocal converts src to a []byte if possible;
+// otherwise it returns an error. If src == nil, ReadSourceLocal returns
+// the result of reading the file specified by filename.
+func ReadSourceLocal(filename string, src any) ([]byte, error) {
+	if src != nil {
+		return ReadSource(src)
+	}
+	return os.ReadFile(filename)
+}
+
+// ReadSourceFromURI reads the source from the given URI.
+// If src != nil, it reads from src; otherwise, it opens the URI and reads
+// from it.
+func ReadSourceFromURI(uri string, src any) (ret []byte, err error) {
+	if src == nil {
+		var f io.ReadCloser
+		f, err = Open(uri)
+		if err != nil {
+			return
+		}
+		defer f.Close()
+		src = f
+	}
+	return ReadSource(src)
+}
+
+// -------------------------------------------------------------------------------------