lib/file/file.go
// Package file defines functions that manipulate files, it's inspired by file helpers from Amoy.
package file
import (
"fmt"
"sync"
dc "github.com/1set/starlet/dataconv"
tps "github.com/1set/starlet/dataconv/types"
"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('file', 'trim_bom')
const ModuleName = "file"
var (
once sync.Once
fileModule starlark.StringDict
none = starlark.None
)
// LoadModule loads the file module. It is concurrency-safe and idempotent.
func LoadModule() (starlark.StringDict, error) {
once.Do(func() {
fileModule = starlark.StringDict{
ModuleName: &starlarkstruct.Module{
Name: ModuleName,
Members: starlark.StringDict{
"copyfile": starlark.NewBuiltin(ModuleName+".copyfile", copyFile),
"stat": starlark.NewBuiltin(ModuleName+".stat", getFileStat),
"trim_bom": starlark.NewBuiltin(ModuleName+".trim_bom", trimBom),
"count_lines": starlark.NewBuiltin(ModuleName+".count_lines", countLinesInFile),
"head_lines": readTopOrBottomLines("head_lines", ReadFirstLines),
"tail_lines": readTopOrBottomLines("tail_lines", ReadLastLines),
"read_bytes": wrapReadFile("read_bytes", readBytes),
"read_string": wrapReadFile("read_string", readString),
"read_lines": wrapReadFile("read_lines", readLines),
"read_json": wrapReadFile("read_json", readJSON),
"read_jsonl": wrapReadFile("read_jsonl", readJSONL),
"write_bytes": wrapWriteFile("write_bytes", true, writeBytes),
"write_string": wrapWriteFile("write_string", true, writeString),
"write_lines": wrapWriteFile("write_lines", true, writeLines),
"write_json": wrapWriteFile("write_json", true, writeJSON),
"write_jsonl": wrapWriteFile("write_jsonl", true, writeJSONL),
"append_bytes": wrapWriteFile("append_bytes", false, writeBytes),
"append_string": wrapWriteFile("append_string", false, writeString),
"append_lines": wrapWriteFile("append_lines", false, writeLines),
"append_json": wrapWriteFile("append_json", false, writeJSON),
"append_jsonl": wrapWriteFile("append_jsonl", false, writeJSONL),
},
},
}
})
return fileModule, nil
}
// trimBom removes the UTF-8 BOM (Byte Order Mark) from the beginning of a string.
func trimBom(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
if l := len(args); l != 1 {
return nil, fmt.Errorf(`%s: takes exactly one argument (%d given)`, b.Name(), l)
}
switch r := args[0]; v := r.(type) {
case starlark.String:
return starlark.String(TrimUTF8BOM([]byte(v))), nil
case starlark.Bytes:
return starlark.Bytes(TrimUTF8BOM([]byte(v))), nil
default:
return none, fmt.Errorf(`%s: expected string or bytes, got %s`, b.Name(), r.Type())
}
}
// countLinesInFile counts the number of lines in a file.
func countLinesInFile(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var fp tps.StringOrBytes
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "name", &fp); err != nil {
return starlark.None, err
}
// do the work
cnt, err := CountFileLines(fp.GoString())
if err != nil {
return nil, err
}
return starlark.MakeInt(cnt), nil
}
// readTopOrBottomLines wraps the file reading functions for top or bottom lines to be used in Starlark.
func readTopOrBottomLines(funcName string, workLoad func(name string, n int) ([]string, error)) starlark.Callable {
return starlark.NewBuiltin(ModuleName+"."+funcName, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
// unpack arguments
var (
fp tps.StringOrBytes
n starlark.Int
)
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "name", &fp, "n", &n); err != nil {
return starlark.None, err
}
nInt, _ := n.Int64()
if nInt <= 0 {
return starlark.None, fmt.Errorf(`%s: expected positive integer, got %d`, b.Name(), n)
}
// read lines
ls, err := workLoad(fp.GoString(), int(nInt))
if err != nil {
return nil, err
}
// return as list
sl := make([]starlark.Value, len(ls))
for i, l := range ls {
sl[i] = starlark.String(l)
}
return starlark.NewList(sl), nil
})
}
// wrapReadFile wraps the file reading functions to be used in Starlark.
func wrapReadFile(funcName string, workLoad func(name string) (starlark.Value, error)) starlark.Callable {
return starlark.NewBuiltin(ModuleName+"."+funcName, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var fp tps.StringOrBytes
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "name", &fp); err != nil {
return starlark.None, err
}
return workLoad(fp.GoString())
})
}
// readBytes reads the whole named file and returns the contents as bytes.
func readBytes(name string) (starlark.Value, error) {
data, err := ReadFileBytes(name)
if err != nil {
return nil, err
}
return starlark.Bytes(data), nil
}
// readString reads the whole named file and returns the contents as string.
func readString(name string) (starlark.Value, error) {
data, err := ReadFileString(name)
if err != nil {
return nil, err
}
return starlark.String(data), nil
}
// readLines reads the whole named file and returns the contents as a list of lines.
func readLines(name string) (starlark.Value, error) {
ls, err := ReadFileLines(name)
if err != nil {
return nil, err
}
sl := make([]starlark.Value, len(ls))
for i, l := range ls {
sl[i] = starlark.String(l)
}
return starlark.NewList(sl), nil
}
// wrapWriteFile wraps the file writing functions to be used in Starlark.
func wrapWriteFile(funcName string, override bool, workLoad func(name, funcName string, override bool, data starlark.Value) error) starlark.Callable {
fullName := ModuleName + "." + funcName
return starlark.NewBuiltin(fullName, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
var (
fp tps.StringOrBytes
data starlark.Value
)
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "name", &fp, "data", &data); err != nil {
return starlark.None, err
}
return starlark.None, workLoad(fp.GoString(), fullName, override, data)
})
}
// writeBytes writes the given data in bytes into a file.
func writeBytes(name, funcName string, override bool, data starlark.Value) error {
wf := AppendFileBytes
if override {
wf = WriteFileBytes
}
// treat starlark.Bytes and starlark.String as the same type
switch v := data.(type) {
case starlark.Bytes:
return wf(name, []byte(v))
case starlark.String:
return wf(name, []byte(v))
default:
return fmt.Errorf(`%s: expected string or bytes, got %s`, funcName, data.Type())
}
}
// writeString writes the given data in string into a file.
func writeString(name, funcName string, override bool, data starlark.Value) error {
wf := AppendFileString
if override {
wf = WriteFileString
}
// treat starlark.Bytes and starlark.String as the same type
switch v := data.(type) {
case starlark.Bytes:
return wf(name, string(v))
case starlark.String:
return wf(name, string(v))
default:
return fmt.Errorf(`%s: expected string or bytes, got %s`, funcName, data.Type())
}
}
// writeLines writes the lines into a file. The data should be a list, a tuple or a set of strings.
func writeLines(name, funcName string, override bool, data starlark.Value) error {
wf := AppendFileLines
if override {
wf = WriteFileLines
}
// handle all types of iterable, and allow string or bytes
switch v := data.(type) {
case starlark.String:
return wf(name, []string{v.GoString()})
case starlark.Bytes:
return wf(name, []string{string(v)})
case *starlark.List:
return wf(name, convIterStrings(v))
case starlark.Tuple:
return wf(name, convIterStrings(v))
case *starlark.Set:
return wf(name, convIterStrings(v))
default:
return fmt.Errorf(`%s: expected list/tuple/set, got %s`, funcName, data.Type())
}
}
func convIterStrings(lst starlark.Iterable) (lines []string) {
iter := lst.Iterate()
defer iter.Done()
var x starlark.Value
for iter.Next(&x) {
lines = append(lines, dc.StarString(x))
}
return
}