1set/starlet

// Package path defines functions that manipulate directories, it's inspired by pathlib module from Mojo.
package path

import (
    "fmt"
    "os"
    "path/filepath"
    "sync"

    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('path', 'join')
const ModuleName = "path"

var (
    once       sync.Once
    pathModule starlark.StringDict
)

// LoadModule loads the path module. It is concurrency-safe and idempotent.
func LoadModule() (starlark.StringDict, error) {
    once.Do(func() {
        pathModule = starlark.StringDict{
            ModuleName: &starlarkstruct.Module{
                Name: ModuleName,
                Members: starlark.StringDict{
                    "abs":     starlark.NewBuiltin(ModuleName+".abs", absPath),
                    "join":    starlark.NewBuiltin(ModuleName+".join", joinPaths),
                    "exists":  wrapExistPath("exists", checkExistPath),
                    "is_file": wrapExistPath("is_file", checkFileExist),
                    "is_dir":  wrapExistPath("is_dir", checkDirExist),
                    "is_link": wrapExistPath("is_link", checkSymlinkExist),
                    "listdir": starlark.NewBuiltin(ModuleName+".listdir", listDirContents),
                    "getcwd":  starlark.NewBuiltin(ModuleName+".getcwd", getCWD),
                    "chdir":   starlark.NewBuiltin(ModuleName+".chdir", changeCWD),
                    "mkdir":   starlark.NewBuiltin(ModuleName+".mkdir", makeDir),
                },
            },
        }
    })
    return pathModule, nil
}

// absPath returns the absolute representation of path.
func absPath(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    var path string
    if err := starlark.UnpackArgs(b.Name(), args, kwargs, "path", &path); err != nil {
        return nil, err
    }
    // get absolute path
    abs, err := filepath.Abs(path)
    if err != nil {
        return nil, err
    }
    return starlark.String(abs), nil
}

// joinPaths joins any number of path elements into a single path.
func joinPaths(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    // check arguments
    if len(args) < 1 {
        return nil, fmt.Errorf("%s: got %d arguments, want at least 1", b.Name(), len(args))
    }
    // unpack arguments
    paths := make([]string, len(args))
    for i, arg := range args {
        s, ok := starlark.AsString(arg)
        if !ok {
            return nil, fmt.Errorf("%s: for parameter path: got %s, want string", b.Name(), arg.Type())
        }
        paths[i] = s
    }
    // join paths
    joined := filepath.Join(paths...)
    return starlark.String(joined), nil
}

// wrapExistPath wraps the existPath function to be used in Starlark with a given function to check if the path exists.
func wrapExistPath(funcName string, workLoad func(path string) bool) starlark.Callable {
    return starlark.NewBuiltin(ModuleName+"."+funcName, func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
        var path string
        if err := starlark.UnpackArgs(b.Name(), args, kwargs, "path", &path); err != nil {
            return starlark.None, err
        }
        return starlark.Bool(workLoad(path)), nil
    })
}

// checkExistPath returns true if the path exists, if it's a symbolic link, the symbolic link is followed.
func checkExistPath(path string) bool {
    _, err := os.Stat(path)
    return err == nil
}

// checkFileExist returns true if the file exists, if it's a symbolic link, the symbolic link is followed.
func checkFileExist(path string) bool {
    info, err := os.Stat(path)
    return err == nil && info != nil && info.Mode().IsRegular()
}

// checkDirExist returns true if the directory exists, if it's a symbolic link, the symbolic link is followed.
func checkDirExist(path string) bool {
    info, err := os.Stat(path)
    return err == nil && info != nil && info.IsDir()
}

// checkSymlinkExist returns true if the symbolic link exists.
func checkSymlinkExist(path string) bool {
    info, err := os.Lstat(path)
    return err == nil && info != nil && info.Mode()&os.ModeSymlink == os.ModeSymlink
}

// listDirContents returns a list of directory contents.
func listDirContents(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    var (
        path       string
        recursive  bool
        filterFunc = tps.NullableCallable{}
    )
    if err := starlark.UnpackArgs(b.Name(), args, kwargs, "path", &path, "recursive?", &recursive, "filter?", &filterFunc); err != nil {
        return nil, err
    }
    // get filter func
    var ff starlark.Callable
    if !filterFunc.IsNull() {
        ff = filterFunc.Value()
    }

    // check root stat
    rootInfo, err := os.Lstat(path)
    if err != nil {
        return nil, fmt.Errorf("%s: %w", b.Name(), err)
    }

    // check if path is a directory, if not return empty list
    var sl []starlark.Value
    if !rootInfo.IsDir() {
        return starlark.NewList(sl), nil
    }

    // scan directory contents
    if err := filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
        if err != nil {
            return err
        }
        // skip same path to avoid infinite loop in case of symbolic links
        if os.SameFile(rootInfo, info) {
            return nil
        }
        // filter path
        sp := starlark.String(p)
        if ff != nil {
            filtered, err := starlark.Call(thread, ff, starlark.Tuple{sp}, nil)
            if err != nil {
                return fmt.Errorf("filter %q: %w", p, err)
            }
            if fb, ok := filtered.(starlark.Bool); !ok {
                return fmt.Errorf("filter %q: got %s, want bool", p, filtered.Type())
            } else if fb == false {
                return nil // skip path
            }
        }

        // add path to list
        sl = append(sl, sp)

        // check if we should list recursively
        if !recursive && p != path && info.IsDir() {
            return filepath.SkipDir
        }
        return nil
    }); err != nil {
        return nil, fmt.Errorf("%s: %w", b.Name(), err)
    }
    return starlark.NewList(sl), nil
}

// getCWD returns the current working directory.
func getCWD(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
    }
    // get current working directory
    cwd, err := os.Getwd()
    if err != nil {
        return nil, fmt.Errorf("%s: %w", b.Name(), err)
    }
    return starlark.String(cwd), nil
}

// changeCWD changes the current working directory.
func changeCWD(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    var path string
    if err := starlark.UnpackArgs(b.Name(), args, kwargs, "path", &path); err != nil {
        return nil, err
    }
    // change working directory
    if err := os.Chdir(path); err != nil {
        return nil, fmt.Errorf("%s: %w", b.Name(), err)
    }
    return starlark.None, nil
}

// makeDir creates a directory with the given name.
func makeDir(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    var (
        pathVal tps.StringOrBytes
        modeVal = uint32(0755)
    )
    if err := starlark.UnpackArgs(b.Name(), args, kwargs, "path", &pathVal, "mode?", &modeVal); err != nil {
        return starlark.None, err
    }
    // do the work
    mode := os.FileMode(modeVal)
    return starlark.None, os.MkdirAll(pathVal.GoString(), mode)
}