run.go
package starlet
import (
"context"
"fmt"
"io/fs"
"sync"
"time"
"github.com/1set/starlet/lib/goidiomatic"
"github.com/1set/starlight/convert"
"go.starlark.net/repl"
"go.starlark.net/starlark"
"go.starlark.net/syntax"
)
// REPL is a Read-Eval-Print-Loop for Starlark.
// It loads the predeclared symbols and modules into the global environment,
func (m *Machine) REPL() {
if err := m.prepareThread(nil); err != nil {
repl.PrintError(err)
return
}
repl.REPLOptions(m.getFileOptions(), m.thread, m.predeclared)
}
// RunScript initiates a Machine, executes a script with extra variables, and returns the Machine and the execution result.
func RunScript(content []byte, extras StringAnyMap) (*Machine, StringAnyMap, error) {
m := NewDefault()
res, err := m.RunScript(content, extras)
return m, res, err
}
// RunFile initiates a Machine, executes a script from a file with extra variables, and returns the Machine and the execution result.
func RunFile(name string, fileSys fs.FS, extras StringAnyMap) (*Machine, StringAnyMap, error) {
m := NewDefault()
res, err := m.RunFile(name, fileSys, extras)
return m, res, err
}
// RunTrustedScript initiates a Machine, executes a script with all builtin modules loaded and extra variables, returns the Machine and the result.
// Use with caution as it allows script access to file system and network.
func RunTrustedScript(content []byte, globals, extras StringAnyMap) (*Machine, StringAnyMap, error) {
m := NewWithBuiltins(globals, nil, nil)
res, err := m.RunScript(content, extras)
return m, res, err
}
// RunTrustedFile initiates a Machine, executes a script from a file with all builtin modules loaded and extra variables, returns the Machine and the result.
// Use with caution as it allows script access to file system and network.
func RunTrustedFile(name string, fileSys fs.FS, globals, extras StringAnyMap) (*Machine, StringAnyMap, error) {
m := NewWithBuiltins(globals, nil, nil)
res, err := m.RunFile(name, fileSys, extras)
return m, res, err
}
// Run executes a preset script and returns the output.
func (m *Machine) Run() (StringAnyMap, error) {
m.mu.Lock()
defer m.mu.Unlock()
return m.runInternal(context.Background(), nil, true)
}
// RunScript executes a script with additional variables, which take precedence over global variables and modules, returns the result.
func (m *Machine) RunScript(content []byte, extras StringAnyMap) (StringAnyMap, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.scriptName = "direct.star"
m.scriptContent = content
m.scriptFS = nil
return m.runInternal(context.Background(), extras, false)
}
// RunFile executes a script from a file with additional variables, which take precedence over global variables and modules, returns the result.
func (m *Machine) RunFile(name string, fileSys fs.FS, extras StringAnyMap) (StringAnyMap, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.scriptName = name
m.scriptContent = nil
m.scriptFS = fileSys
return m.runInternal(context.Background(), extras, true)
}
// RunWithTimeout executes a preset script with a timeout and additional variables, which take precedence over global variables and modules, returns the result.
func (m *Machine) RunWithTimeout(timeout time.Duration, extras StringAnyMap) (StringAnyMap, error) {
m.mu.Lock()
defer m.mu.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
return m.runInternal(ctx, extras, true)
}
// RunWithContext executes a preset script within a specified context and additional variables, which take precedence over global variables and modules, returns the result.
func (m *Machine) RunWithContext(ctx context.Context, extras StringAnyMap) (StringAnyMap, error) {
m.mu.Lock()
defer m.mu.Unlock()
return m.runInternal(ctx, extras, true)
}
func (m *Machine) runInternal(ctx context.Context, extras StringAnyMap, allowCache bool) (out StringAnyMap, err error) {
defer func() {
if r := recover(); r != nil {
err = errorStarlarkPanic("exec", r)
}
}()
// either script content or name and FS must be set
var (
scriptName = m.scriptName
source interface{}
)
if m.scriptContent != nil {
if scriptName == "" {
// for default name, and disable cache to avoid conflict
scriptName = "eval.star"
allowCache = false
}
source = m.scriptContent
} else if m.scriptFS != nil {
if scriptName == "" {
// if no name, cannot load
return nil, errorStarletErrorf("run", "no script name")
}
// load script from FS
rd, e := m.scriptFS.Open(scriptName)
if e != nil {
return nil, errorStarletError("run", e)
}
source = rd
} else {
return nil, errorStarletErrorf("run", "no script to execute")
}
// prepare thread
if err = m.prepareThread(extras); err != nil {
return nil, err
}
// cancel thread when context cancelled
if ctx == nil || ctx.Err() != nil {
// for nil context, or context already cancelled, use a new one
ctx = context.TODO()
}
m.thread.SetLocal("context", ctx)
// wait for the routine to finish, or cancel it when context cancelled
var wg sync.WaitGroup
defer wg.Wait()
wg.Add(1)
// if context is not cancelled, cancel the routine when execution is done, or panic
done := make(chan struct{}, 1)
defer close(done)
go func() {
defer wg.Done()
select {
case <-ctx.Done():
m.thread.Cancel("context cancelled")
case <-done:
// No action if the script has finished
}
}()
// run with everything prepared
m.runTimes++
res, err := m.execStarlarkFile(scriptName, source, allowCache)
done <- struct{}{}
// merge result as predeclared for next run
for k, v := range res {
m.predeclared[k] = v
}
// handle result and convert
out = m.convertOutput(res)
if err != nil {
// for exit code
if err.Error() == goidiomatic.ErrSystemExit.Error() {
var exitCode uint8
if c := m.thread.Local("exit_code"); c != nil {
if co, ok := c.(uint8); ok {
exitCode = co
}
}
// exit code 0 means success
if exitCode == 0 {
err = nil
} else {
err = errorStarletErrorf("run", "exit code: %d", exitCode)
}
} else {
// wrap starlark errors
err = errorStarlarkError("exec", err)
}
return out, err
}
return out, nil
}
// prepareThread prepares the thread for execution, including preset globals, preload modules and extras.
func (m *Machine) prepareThread(extras StringAnyMap) (err error) {
mergeExtra := func() error {
// no extras
if extras == nil {
return nil
}
// convert extras if needed
esd, err := m.convertInput(extras)
if err != nil {
return errorStarlightConvert("extras", err)
}
// merge extras
for k, v := range esd {
m.predeclared[k] = v
}
return nil
}
// initialize thread or reset for each run
if m.thread == nil {
// -- for the first run
// preset globals + preload modules + extras -> predeclared
if m.predeclared, err = m.convertInput(m.globals); err != nil {
return errorStarlightConvert("globals", err)
}
if err = m.preloadMods.LoadAll(m.predeclared); err != nil {
return errorStarletError("preload", err)
}
// merge extras into predeclared
if err = mergeExtra(); err != nil {
return err
}
// cache load&read + printf -> thread
m.loadCache = &cache{
cache: make(map[string]*entry),
loadMod: m.lazyloadMods.GetLazyLoader(),
readFile: func(name string) ([]byte, error) {
return readScriptFile(name, m.scriptFS)
},
globals: m.predeclared,
}
m.thread = &starlark.Thread{
Name: "starlet",
Print: m.printFunc,
Load: func(thread *starlark.Thread, module string) (starlark.StringDict, error) {
return m.loadCache.Load(module)
},
}
} else {
// -- for the second and following runs
// merge extras into predeclared
if err = mergeExtra(); err != nil {
return err
}
// set globals for cache
m.loadCache.loadMod = m.lazyloadMods.GetLazyLoader()
m.loadCache.globals = m.predeclared
// reset for each run
m.thread.Print = m.printFunc
m.thread.Uncancel()
}
return nil
}
// Reset resets the machine to initial state before the first run.
// Attention: It does not reset the compiled program cache.
func (m *Machine) Reset() {
m.runTimes = 0
m.thread = nil
m.loadCache = nil
m.predeclared = nil
}
// convertInput converts a StringAnyMap to a starlark.StringDict, usually for output variable.
func (m *Machine) convertInput(a StringAnyMap) (starlark.StringDict, error) {
if m.enableInConv {
return convert.MakeStringDictWithTag(a, m.customTag)
}
return castStringAnyMapToStringDict(a)
}
// convertOutput converts a starlark.StringDict to a StringAnyMap, usually for output variable.
func (m *Machine) convertOutput(d starlark.StringDict) StringAnyMap {
if m.enableOutConv {
return convert.FromStringDict(d)
}
return castStringDictToAnyMap(d)
}
// getFileOptions gets the exec options from the config.
func (m *Machine) getFileOptions() *syntax.FileOptions {
opt := syntax.FileOptions{
Set: true,
}
if m.allowRecursion {
opt.Recursion = true
}
if m.allowGlobalReassign {
opt.GlobalReassign = true
opt.TopLevelControl = true
opt.While = true
}
return &opt
}
// castStringDictToAnyMap converts a starlark.StringDict to a StringAnyMap without any Starlight conversion.
func castStringDictToAnyMap(m starlark.StringDict) StringAnyMap {
ret := make(StringAnyMap, len(m))
for k, v := range m {
ret[k] = v
}
return ret
}
// castStringAnyMapToStringDict converts a StringAnyMap to a starlark.StringDict without any Starlight conversion.
// It fails if any values are not starlark.Value.
func castStringAnyMapToStringDict(m StringAnyMap) (starlark.StringDict, error) {
ret := make(starlark.StringDict, len(m))
for k, v := range m {
sv, ok := v.(starlark.Value)
if !ok {
return nil, fmt.Errorf("value of key %q is not a starlark.Value", k)
}
ret[k] = sv
}
return ret, nil
}