idiomatic.go

// Package goidiomatic provides a Starlark module that defines Go idiomatic functions and values.
package goidiomatic

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/1set/starlet/dataconv"
	tps "github.com/1set/starlet/dataconv/types"
	"github.com/1set/starlight/convert"
	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"
)

// ModuleName defines the expected name for this Module when used in Starlark's load() function, eg: load('go_idiomatic', 'nil')
const ModuleName = "go_idiomatic"

// LoadModule loads the Go idiomatic module.
func LoadModule() (starlark.StringDict, error) {
	return starlark.StringDict{
		"true":             starlark.True,
		"false":            starlark.False,
		"nil":              starlark.None,
		"length":           starlark.NewBuiltin("length", length),
		"sum":              starlark.NewBuiltin("sum", sum),
		"oct":              starlark.NewBuiltin("oct", oct),
		"hex":              starlark.NewBuiltin("hex", hex),
		"bytes_hex":        starlark.NewBuiltin("bytes_hex", bytesToHex),
		"is_nil":           starlark.NewBuiltin("is_nil", isNil),
		"bin":              starlark.NewBuiltin("bin", bin),
		"sleep":            starlark.NewBuiltin("sleep", sleep),
		"exit":             starlark.NewBuiltin("exit", exit),
		"quit":             starlark.NewBuiltin("quit", exit), // alias for exit
		"module":           starlark.NewBuiltin("module", starlarkstruct.MakeModule),
		"struct":           starlark.NewBuiltin("struct", starlarkstruct.Make),
		"make_struct":      starlark.NewBuiltin("make_struct", makeCustomStruct),
		"shared_dict":      starlark.NewBuiltin("shared_dict", makeSharedDict),
		"make_shared_dict": starlark.NewBuiltin("make_shared_dict", makeCustomSharedDict),
		"to_dict":          starlark.NewBuiltin("to_dict", convertToDict),
		"distinct":         starlark.NewBuiltin("distinct", distinct),
		"eprint":           starlark.NewBuiltin("eprint", stderrPrint),
		"pprint":           starlark.NewBuiltin("pprint", prettyPrint),
	}, nil
}

// for convenience
var (
	none = starlark.None
)

// isNil returns true if the given value is nil or wraps nil inside.
func isNil(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var x starlark.Value
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "x", &x); err != nil {
		return none, err
	}
	switch t := x.(type) {
	// Starlark's NoneType is the only type that can be treated as nil.
	case starlark.NoneType:
		return starlark.True, nil
	// Go types that wrap nil inside.
	case *convert.GoSlice:
		return starlark.Bool(dataconv.IsInterfaceNil(t) || dataconv.IsInterfaceNil(t.Value().Interface())), nil
	case *convert.GoMap:
		return starlark.Bool(dataconv.IsInterfaceNil(t) || dataconv.IsInterfaceNil(t.Value().Interface())), nil
	case *convert.GoStruct:
		return starlark.Bool(dataconv.IsInterfaceNil(t) || dataconv.IsInterfaceNil(t.Value().Interface())), nil
	case *convert.GoInterface:
		return starlark.Bool(dataconv.IsInterfaceNil(t) || dataconv.IsInterfaceNil(t.Value().Interface())), nil
	default:
		return none, fmt.Errorf("%s: unsupported type: %T", b.Name(), t)
	}
}

// distinct returns a iterable with distinct elements from the given iterable, i.e. without duplicates.
// for list and custom types, it returns a new list with distinct elements.
// for tuple, it returns a new tuple with distinct elements.
// for dict, it calls the keys() method to returns the keys in a list.
// for set, it just returns the original set.
func distinct(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var itr starlark.Iterable
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "iterable", &itr); err != nil {
		return none, err
	}
	// check if the iterable is a set or a map, just return the keys quickly
	switch v := itr.(type) {
	case *starlark.Dict: // get the keys of the dict
		return starlark.NewList(v.Keys()), nil
	case *starlark.Set: // get the original set
		return v, nil
	}
	// get the list elements for list or tuple
	var (
		lsv  []starlark.Value
		hm   = make(map[uint32]struct{})
		x    starlark.Value
		iter = itr.Iterate()
	)
	// loop through the list/tuple and add distinct elements
	for iter.Next(&x) {
		h, e := x.Hash()
		if e != nil {
			return none, fmt.Errorf("%s: %w", b.Name(), e)
		}
		if _, ok := hm[h]; !ok {
			hm[h] = struct{}{}
			lsv = append(lsv, x)
		}
	}
	// return the new iterable as per the input type
	switch itr.(type) {
	case *starlark.List:
		return starlark.NewList(lsv), nil
	case starlark.Tuple:
		return starlark.Tuple(lsv), nil
	default:
		// fallback to list for other types, like the custom ones
		return starlark.NewList(lsv), nil
	}
}

// length returns the length of the given value.
func length(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	if l := len(args); l != 1 {
		return none, fmt.Errorf(`length() takes exactly one argument (%d given)`, l)
	}

	switch r := args[0]; v := r.(type) {
	case starlark.String:
		return starlark.MakeInt(utf8.RuneCountInString(v.GoString())), nil
	case starlark.Bytes:
		return starlark.MakeInt(len(v)), nil
	default:
		if sv, ok := v.(starlark.Sequence); ok {
			return starlark.MakeInt(sv.Len()), nil
		}
		return none, fmt.Errorf(`length() function isn't supported for '%s' type object`, v.Type())
	}
}

// sum returns the sum of the given values.
func sum(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var (
		lst   starlark.Iterable
		total = tps.NewNumericValue()
	)
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "iterable", &lst, "start?", total); err != nil {
		return none, err
	}

	// loop through the list
	iter := lst.Iterate()
	defer iter.Done()

	var x starlark.Value
	for iter.Next(&x) {
		if err := total.Add(x); err != nil {
			return none, err
		}
	}

	// return the result
	return total.Value(), nil
}

// bytesToHex returns the hexadecimal representation of the given bytes.
func bytesToHex(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	// parse and set the arguments
	var (
		bs  starlark.Bytes
		sep starlark.String
		bps = starlark.MakeInt(1)
	)
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "bytes", &bs, "sep?", &sep, "bytes_per_sep?", &bps); err != nil {
		return none, err
	}
	bytesPerSep, ok := bps.Int64()
	if !ok {
		return none, fmt.Errorf("invalid bytes_per_sep: %v", bps)
	}
	// convert the bytes to hexadecimal, positive values calculate the separator position from the right, negative values from the left.
	var (
		bpsInt   = int(bytesPerSep)
		sepStr   = string(sep)
		hexBytes = make([]string, len(bs))
	)
	for i, b := range bs {
		hexBytes[i] = fmt.Sprintf("%02x", b)
	}
	if bpsInt > 0 { // for downto
		for i := len(hexBytes) - bpsInt; i > 0; i -= bpsInt {
			hexBytes[i] = sepStr + hexBytes[i]
		}
	} else if bpsInt < 0 { // for to
		for i := -bpsInt; i < len(hexBytes); i -= bpsInt {
			hexBytes[i] = sepStr + hexBytes[i]
		}
	}
	// compile the result
	var rs strings.Builder
	rs.WriteString(strings.Join(hexBytes, ""))
	return starlark.String(rs.String()), nil
}

// hex returns the hexadecimal representation of the given value.
func hex(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var x starlark.Int
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "x", &x); err != nil {
		return none, err
	}
	return convertStarlarkNumber(x, 16, "0x")
}

// oct returns the octal representation of the given value.
func oct(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var x starlark.Int
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "x", &x); err != nil {
		return none, err
	}
	return convertStarlarkNumber(x, 8, "0o")
}

// bin returns the binary representation of the given value.
func bin(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var x starlark.Int
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "x", &x); err != nil {
		return none, err
	}
	return convertStarlarkNumber(x, 2, "0b")
}

// convertStarlarkNumber converts the given starlark.Int number to the given base string.
func convertStarlarkNumber(x starlark.Int, base int, fmtPre string) (starlark.Value, error) {
	var (
		s    string
		n    = x.BigInt()
		sign = n.Sign()
	)
	if sign == 0 {
		s = fmtPre + "0"
	} else {
		signPre := ""
		if sign < 0 {
			signPre = "-"
			n.Neg(n)
		}
		s = signPre + fmtPre + n.Text(base)
	}
	return starlark.String(s), nil
}

// sleep sleeps for the given number of seconds.
func sleep(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	// get the duration
	var sec tps.FloatOrInt
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "secs", &sec); err != nil {
		return none, err
	}
	if sec < 0 {
		return none, errors.New("secs must be non-negative")
	}
	dur := time.Duration(float64(sec) * float64(time.Second))
	// get the context
	ctx := dataconv.GetThreadContext(thread)
	// sleep
	t := time.NewTimer(dur)
	defer t.Stop()
	select {
	case <-t.C:
		return none, nil
	case <-ctx.Done():
		return none, ctx.Err()
	}
}

// exit exits the program with the given exit code.
func exit(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var code uint8
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "code?", &code); err != nil {
		return none, err
	}
	thread.SetLocal("exit_code", code)
	return none, ErrSystemExit
}

var (
	// ErrSystemExit is returned by exit() to indicate the program should exit.
	ErrSystemExit = errors.New(`starlet runtime system exit (Use Ctrl-D in REPL to exit)`)
)

// makeCustomStruct creates a custom struct with the given name and members.
func makeCustomStruct(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var ctor starlark.Value
	if err := starlark.UnpackPositionalArgs(b.Name(), args, nil, 1, &ctor); err != nil {
		return nil, err
	}
	return starlarkstruct.FromKeywords(ctor, kwargs), nil
}

// makeSharedDict creates a new shared dictionary with default name.
func makeSharedDict(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	// check the arguments: no arguments
	if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil {
		return nil, err
	}
	// do the job
	return dataconv.NewSharedDict(), nil
}

// makeCustomSharedDict creates a custom shared dictionary with the given name and default members.
func makeCustomSharedDict(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var (
		name string
		data *starlark.Dict
	)
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "name?", &name, "data?", &data); err != nil {
		return nil, err
	}
	// create the shared data
	if name == "" && data == nil {
		// if no name or data given, create the default shared data
		return dataconv.NewSharedDict(), nil
	} else if data == nil {
		// if only name given, create the named shared data
		return dataconv.NewNamedSharedDict(name), nil
	} else {
		// if both name and data given, create the named shared data with data
		sd := dataconv.NewSharedDictFromDict(data)
		sd.SetTypeName(name)
		return sd, nil
	}
}

// stderrPrint works like standard print() but prints the given arguments to stderr.
func stderrPrint(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	sep := " "
	if err := starlark.UnpackArgs(b.Name(), nil, kwargs, "sep?", &sep); err != nil {
		return nil, err
	}
	// build output string
	buf := new(strings.Builder)
	for i, v := range args {
		if i > 0 {
			buf.WriteString(sep)
		}
		// convert to string
		buf.WriteString(dataconv.StarString(v))
	}
	// write to stderr
	s := buf.String()
	fmt.Fprintln(os.Stderr, s)
	return starlark.None, nil
}

// prettyPrint works like standard print() but formats the given arguments in pretty JSON format with indentation.
func prettyPrint(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	sep := " "
	if err := starlark.UnpackArgs(b.Name(), nil, kwargs, "sep?", &sep); err != nil {
		return nil, err
	}
	// build output string
	buf := new(strings.Builder)
	for i, v := range args {
		if i > 0 {
			buf.WriteString(sep)
		}
		// convert to JSON or string as fallback
		raw, err := dataconv.MarshalStarlarkJSON(v, 4)
		if err != nil {
			buf.WriteString(dataconv.StarString(v))
		} else {
			buf.WriteString(raw)
		}
	}
	// write like std print
	s := buf.String()
	if thread.Print != nil {
		thread.Print(thread, s)
	} else {
		fmt.Fprintln(os.Stderr, s)
	}
	return starlark.None, nil
}

// convertToDict creates a Starlark dict from a Starlark dict, module, struct, GoStruct, or SharedDict.
// It works as a complement to the builtin dict() function of Starlark, not a replacement or alternative.
// For GoStruct, it converts the underlying Go struct to JSON string using Go standard JSON encoder, and then decodes it to a Starlark dict with Starlark JSON decoder.
// Errors are returned if the conversion fails.
func convertToDict(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var v starlark.Value
	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "v", &v); err != nil {
		return nil, err
	}
	// convert to dict
	switch t := v.(type) {
	case *starlarkstruct.Module:
		dt := starlark.NewDict(len(t.Members))
		for k, v := range t.Members {
			_ = dt.SetKey(starlark.String(k), v)
		}
		return dt, nil
	case *starlarkstruct.Struct:
		sd := starlark.StringDict{}
		t.ToStringDict(sd)
		dt := starlark.NewDict(len(sd))
		for k, v := range sd {
			_ = dt.SetKey(starlark.String(k), v)
		}
		return dt, nil
	case *convert.GoStruct:
		rv := t.Value().Interface()
		bs, err := json.Marshal(rv)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", b.Name(), err)
		}
		return dataconv.DecodeStarlarkJSON(bs)
	case *dataconv.SharedDict:
		return t.CloneDict()
	case *starlark.Dict:
		return dataconv.CloneDict(t)
	default:
		return nil, fmt.Errorf("%s: unsupported type: %T", b.Name(), t)
	}
}