Merge pull request #140 from xushiwei/q

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,78 @@ 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) ([]byte, error) {
+	if src == nil {
+		f, err := Open(uri)
+		if err != nil {
+			return nil, err
+		}
+		defer f.Close()
+		src = f
+	}
+	return ReadSource(src)
+}
+
+// -------------------------------------------------------------------------------------

stream/stream_test.go

@@ -17,7 +17,9 @@
 package stream_test
 
 import (
+	"bytes"
 	"io"
+	"strings"
 	"testing"
 
 	"github.com/qiniu/x/stream"
@@ -51,3 +53,38 @@ func TestOpenFile(t *testing.T) {
 		t.Fatal("Open local file success?")
 	}
 }
+
+func TestReadSource(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	if _, err := stream.ReadSource(buf); err != nil {
+		t.Fatal("readSource failed:", err)
+	}
+	sr := strings.NewReader("")
+	if _, err := stream.ReadSource(sr); err != nil {
+		t.Fatal("readSource strings.Reader failed:", err)
+	}
+	if _, err := stream.ReadSource(0); err == nil {
+		t.Fatal("readSource int failed: no error?")
+	}
+}
+
+func TestReadSourceLocal(t *testing.T) {
+	if _, err := stream.ReadSourceLocal("", ""); err != nil {
+		t.Fatal("ReadSourceLocal failed:", err)
+	}
+	if _, err := stream.ReadSourceLocal("/foo/bar/not-exists", nil); err == nil {
+		t.Fatal("ReadSourceLocal int failed: no error?")
+	}
+}
+
+func TestReadSourceFromURI(t *testing.T) {
+	if _, err := stream.ReadSourceFromURI("", []byte("abc")); err != nil {
+		t.Fatal("ReadSourceFromURI failed:", err)
+	}
+	if _, err := stream.ReadSourceFromURI("inline:abc", nil); err != nil {
+		t.Fatal("ReadSourceFromURI failed:", err)
+	}
+	if _, err := stream.ReadSourceFromURI("unknown:abc", nil); err == nil {
+		t.Fatal("ReadSourceFromURI failed: no error?")
+	}
+}