-
Notifications
You must be signed in to change notification settings - Fork 111
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
runtime: Add the ability to read files from Starlark
Starlark scripts can now load arbitrary files from their bundle and read them using a Python-like `open()` and `read()` API.
- Loading branch information
1 parent
ce0d274
commit 08d2939
Showing
4 changed files
with
227 additions
and
40 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,106 @@ | ||
| package runtime | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "io" | ||
| "io/fs" | ||
|
|
||
| "go.starlark.net/starlark" | ||
| "go.starlark.net/starlarkstruct" | ||
| "tidbyt.dev/pixlet/starlarkutil" | ||
| ) | ||
|
|
||
| type File struct { | ||
| fsys fs.FS | ||
| path string | ||
| } | ||
|
|
||
| func (f File) Struct() *starlarkstruct.Struct { | ||
| return starlarkstruct.FromStringDict(starlark.String("File"), starlark.StringDict{ | ||
| "path": starlark.String(f.path), | ||
| "open": starlark.NewBuiltin("open", f.open), | ||
| }) | ||
| } | ||
|
|
||
| func (f File) open(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { | ||
| var mode starlark.String | ||
| if err := starlark.UnpackArgs("open", args, kwargs, "mode?", &mode); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| var binaryMode bool | ||
| switch mode.GoString() { | ||
| case "", "r", "rt": | ||
| binaryMode = false | ||
|
|
||
| case "rb": | ||
| binaryMode = true | ||
|
|
||
| default: | ||
| return nil, fmt.Errorf("unsupported mode: %s", mode) | ||
| } | ||
|
|
||
| fl, err := f.fsys.Open(f.path) | ||
| if err != nil { | ||
| return nil, err | ||
| } else { | ||
| starlarkutil.AddOnExit(thread, func() { fl.Close() }) | ||
| } | ||
|
|
||
| return Reader{fl, binaryMode}.Struct(), nil | ||
| } | ||
|
|
||
| type Reader struct { | ||
| io.ReadCloser | ||
| binaryMode bool | ||
| } | ||
|
|
||
| func (r Reader) Struct() *starlarkstruct.Struct { | ||
| return starlarkstruct.FromStringDict(starlark.String("Reader"), starlark.StringDict{ | ||
| "read": starlark.NewBuiltin("read", r.read), | ||
| "close": starlark.NewBuiltin("close", func(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { | ||
| return nil, r.Close() | ||
| }), | ||
| }) | ||
| } | ||
|
|
||
| // read reads the contents of the file. The Starlark signature is: | ||
| // | ||
| // read(size=-1) -> bytes | ||
| func (r Reader) read(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { | ||
| starlarkSize := starlark.MakeInt(-1) | ||
| if err := starlark.UnpackArgs("read", args, kwargs, "size?", &starlarkSize); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| var size int | ||
| if err := starlark.AsInt(starlarkSize, &size); err != nil { | ||
| return nil, fmt.Errorf("size is not an int") | ||
| } | ||
|
|
||
| returnType := func(buf []byte) starlark.Value { | ||
| if r.binaryMode { | ||
| return starlark.Bytes(buf) | ||
| } else { | ||
| return starlark.String(buf) | ||
| } | ||
| } | ||
|
|
||
| if size < 0 { | ||
| // read and return all bytes | ||
| buf, err := io.ReadAll(r) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| return returnType(buf), nil | ||
| } else { | ||
| // read and return size bytes | ||
| buf := make([]byte, size) | ||
| _, err := r.Read(buf) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return returnType(buf), nil | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,41 @@ | ||
| package runtime | ||
|
|
||
| import ( | ||
| "testing" | ||
| "testing/fstest" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestReadFile(t *testing.T) { | ||
| src := ` | ||
| load("hello.txt", hello = "file") | ||
| def assert_eq(message, actual, expected): | ||
| if not expected == actual: | ||
| fail(message, "-", "expected", expected, "actual", actual) | ||
| def test_read(): | ||
| f = hello.open() | ||
| assert_eq("read", f.read(), "hello world") | ||
| def test_read_binary(): | ||
| f = hello.open(mode="rb") | ||
| assert_eq("read", f.read(), b"hello world") | ||
| def main(): | ||
| pass | ||
| ` | ||
|
|
||
| helloTxt := `hello world` | ||
|
|
||
| vfs := &fstest.MapFS{ | ||
| "main.star": {Data: []byte(src)}, | ||
| "hello.txt": {Data: []byte(helloTxt)}, | ||
| } | ||
|
|
||
| app, err := NewAppletFromFS("test_read_file", vfs) | ||
| require.NoError(t, err) | ||
| app.RunTests(t) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| package starlarkutil | ||
|
|
||
| import "go.starlark.net/starlark" | ||
|
|
||
| const ( | ||
| // ThreadOnExitKey is the key used to store functions that should be called | ||
| // when a thread exits. | ||
| ThreadOnExitKey = "tidbyt.dev/pixlet/runtime/on_exit" | ||
| ) | ||
|
|
||
| type threadOnExitFunc func() | ||
|
|
||
| func AddOnExit(thread *starlark.Thread, fn threadOnExitFunc) { | ||
| if onExit, ok := thread.Local(ThreadOnExitKey).(*[]threadOnExitFunc); ok { | ||
| *onExit = append(*onExit, fn) | ||
| } else { | ||
| thread.SetLocal(ThreadOnExitKey, &[]threadOnExitFunc{fn}) | ||
| } | ||
| } | ||
|
|
||
| func RunOnExitFuncs(thread *starlark.Thread) { | ||
| if onExit, ok := thread.Local(ThreadOnExitKey).(*[]threadOnExitFunc); ok { | ||
| for _, fn := range *onExit { | ||
| fn() | ||
| } | ||
| } | ||
| } |