lure-updater/internal/builtins/reader.go

239 lines
6.2 KiB
Go

package builtins
import (
"bufio"
"encoding/json"
"errors"
"io"
"strings"
"github.com/vmihailenco/msgpack/v5"
"go.elara.ws/lure-updater/internal/convert"
"go.starlark.net/starlark"
"go.starlark.net/starlarkstruct"
)
type starlarkReader struct {
closeFunc func() error
br *bufio.Reader
*starlarkstruct.Struct
}
func newStarlarkReader(r io.Reader) starlarkReader {
sr := starlarkReader{br: bufio.NewReader(r)}
if rc, ok := r.(io.ReadCloser); ok {
sr.closeFunc = rc.Close
}
sr.Struct = starlarkstruct.FromStringDict(starlark.String("regex"), starlark.StringDict{
"read": starlark.NewBuiltin("reader.read", sr.read),
"peek": starlark.NewBuiltin("reader.peek", sr.peek),
"discard": starlark.NewBuiltin("reader.discard", sr.discard),
"read_string": starlark.NewBuiltin("reader.read_string", sr.readString),
"read_until": starlark.NewBuiltin("reader.read_until", sr.readUntil),
"read_string_until": starlark.NewBuiltin("reader.read_string_until", sr.readStringUntil),
"read_all": starlark.NewBuiltin("reader.read_all", sr.readAll),
"read_all_string": starlark.NewBuiltin("reader.read_all_string", sr.readAllString),
"read_json": starlark.NewBuiltin("reader.read_json", sr.readJSON),
"read_msgpack": starlark.NewBuiltin("reader.read_msgpack", sr.readMsgpack),
"close": starlark.NewBuiltin("reader.close", sr.closeReader),
})
return sr
}
func (sr starlarkReader) read(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var n int
err := starlark.UnpackArgs("reader.read", args, kwargs, "n", &n)
if err != nil {
return nil, err
}
buf := make([]byte, n)
_, err = io.ReadFull(sr.br, buf)
if err != nil {
return nil, err
}
return starlark.Bytes(buf), nil
}
func (sr starlarkReader) readString(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var n int
err := starlark.UnpackArgs("reader.read_string", args, kwargs, "n", &n)
if err != nil {
return nil, err
}
buf := make([]byte, n)
_, err = io.ReadFull(sr.br, buf)
if err != nil {
return nil, err
}
return starlark.String(buf), nil
}
func (sr starlarkReader) readUntil(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var delimiter string
err := starlark.UnpackArgs("reader.read_until", args, kwargs, "delimiter", &delimiter)
if err != nil {
return nil, err
}
buf, err := sr.br.ReadBytes(delimiter[0])
if err != nil {
return nil, err
}
return starlark.Bytes(buf), nil
}
func (sr starlarkReader) readStringUntil(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var delimiter string
err := starlark.UnpackArgs("reader.read_string_until", args, kwargs, "delimiter", &delimiter)
if err != nil {
return nil, err
}
buf, err := sr.br.ReadString(delimiter[0])
if err != nil {
return nil, err
}
return starlark.String(buf), nil
}
func (sr starlarkReader) peek(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var n int
err := starlark.UnpackArgs("reader.peek", args, kwargs, "n", &n)
if err != nil {
return nil, err
}
buf, err := sr.br.Peek(n)
if err != nil {
return nil, err
}
return starlark.Bytes(buf), nil
}
func (sr starlarkReader) discard(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var n int
err := starlark.UnpackArgs("reader.discard", args, kwargs, "n", &n)
if err != nil {
return nil, err
}
dn, err := sr.br.Discard(n)
if err != nil {
return nil, err
}
return starlark.MakeInt(dn), nil
}
func (sr starlarkReader) readAll(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var limit int64 = 102400
err := starlark.UnpackArgs("reader.read_all", args, kwargs, "limit??", &limit)
if err != nil {
return nil, err
}
var r io.Reader = sr.br
if limit > 0 {
r = io.LimitReader(sr.br, limit)
}
buf, err := io.ReadAll(r)
if err != nil {
return nil, err
}
return starlark.Bytes(buf), nil
}
func (sr starlarkReader) readAllString(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var limit int64 = 102400
err := starlark.UnpackArgs("reader.read_all_string", args, kwargs, "limit??", &limit)
if err != nil {
return nil, err
}
var r io.Reader = sr.br
if limit > 0 {
r = io.LimitReader(sr.br, limit)
}
buf, err := io.ReadAll(r)
if err != nil {
return nil, err
}
return starlark.String(buf), nil
}
func (sr starlarkReader) readJSON(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var v any
err := json.NewDecoder(sr.br).Decode(&v)
if err != nil {
return nil, err
}
return convert.Convert(v)
}
func (sr starlarkReader) readMsgpack(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var v any
err := msgpack.NewDecoder(sr.br).Decode(&v)
if err != nil {
return nil, err
}
return convert.Convert(v)
}
func (sr starlarkReader) closeReader(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
if sr.closeFunc != nil {
err := sr.closeFunc()
if err != nil {
return nil, err
}
}
return starlark.None, nil
}
// Read implements the io.ReadCloser interface
func (sr starlarkReader) Read(b []byte) (int, error) {
return sr.br.Read(b)
}
// Close implements the io.ReadCloser interface
func (sr starlarkReader) Close() error {
if sr.closeFunc != nil {
return sr.closeFunc()
}
return nil
}
type readerValue struct {
io.ReadCloser
}
func (rv *readerValue) Unpack(v starlark.Value) error {
switch val := v.(type) {
case starlark.String:
rv.ReadCloser = io.NopCloser(strings.NewReader(string(val)))
case starlark.Bytes:
rv.ReadCloser = io.NopCloser(strings.NewReader(string(val)))
case starlarkReader:
rv.ReadCloser = val
}
if rv.ReadCloser == nil {
return errors.New("invalid type for reader")
}
return nil
}