gopp/tengo.go

package tpp

import (
	"github.com/d5/tengo/v2"
	"fmt"
	"context"
	"bytes"
	"strconv"
)

type CompiledMap map[string] *Compiled

type Script = tengo.Script
type Compiled = tengo.Compiled

// The type describes
// customizable way to evaluate
// files.
type Tengo struct {
	// Functions to modify
	// preprocessor for
	// more specific purposes.
	preCode func(context.Context) []byte
	postCode func(context.Context) []byte
	preCompile func(context.Context, *Script)
}

// Returns the new Tengo preprocessor
// with 
func NewTengo() *Tengo {
	ret := &Tengo{}
	return ret
}

func (tengo *Tengo) SetPreCode(fn func(context.Context) []byte) *Tengo {
	tengo.preCode = fn
	return tengo
}

func (tengo *Tengo) SetPostCode(fn func(context.Context) []byte) *Tengo {
	tengo.postCode = fn
	return tengo
}

func (tengo *Tengo) SetPreCompile(fn func(context.Context, *Script)) *Tengo {
	tengo.preCompile = fn
	return tengo
}

func (pp *Tengo) Eval(
	ctx context.Context,
	// File path to give it.
	filePath string,
	// Static text pieces.
	texts [][]byte,
	// The code that comes before any
	// dynamic code.
	preCode []byte,
	// General code.
	codes [][]byte,
) (*Compiled, []byte, error) {
	var fullCodeBuf, retBuf bytes.Buffer
	const retHead = `
		context := {}
		pp := immutable({
			filepath: __pp_filepath__,
			printf : __pp_printf__,
			print : __pp_print__,
			println : __pp_println__,
			write_raw : __pp_write_raw__
		})
	`

	const retSeparator = `
		__separate__()
	`

	fmt.Fprint(&fullCodeBuf, retHead)

	if preCode != nil {
		fmt.Fprintln(
			&fullCodeBuf,
			string(preCode)+"\n",
		)
	}

	if pp.preCode != nil {
		fullCodeBuf.Write(pp.preCode(ctx))
		fullCodeBuf.Write([]byte(retSeparator))
	}

	for i, code := range codes {
		fmt.Fprintf(
			&fullCodeBuf,
			"\n__pp_write_raw__(%q)\n",
			texts[i],
		)
		fmt.Fprintln(
			&fullCodeBuf,
			"\n" + string(code) + retSeparator,
		)
	}
	fmt.Fprintf(
		&fullCodeBuf,
		"\n__pp_write_raw__(%q)\n",
		texts[len(texts)-1],
	)

	if pp.postCode != nil {
		fullCodeBuf.Write(pp.postCode(ctx))
		fullCodeBuf.Write([]byte(retSeparator))
	}

	script := tengo.NewScript(fullCodeBuf.Bytes())
	if pp.preCompile != nil {
		pp.preCompile(ctx, script)
	}

	script.Add("__pp_filepath__", filePath)
	script.Add("__pp_printf__", &tengo.UserFunction{
		Value: func(args ...tengo.Object) (tengo.Object, error){
			if len(args) < 1 {
				return nil, tengo.ErrWrongNumArguments
			}
			format, ok := tengo.ToString(args[0])
			if !ok {
				return nil, tengo.ErrInvalidArgumentType{
					Expected: "string",
				}
			}
			gargs := make([]any, len(args) - 1)
			for i := range gargs {
				gargs[i] = tengo.ToInterface(args[i+1])
				//fmt.Printf("shit: %q\n", gargs[i])
			}
			fmt.Fprintf(&retBuf, format, gargs...)
			return nil, nil
			//return tengo.FromInterface([]byte(str))
		},
	})
	script.Add("__pp_print__", &tengo.UserFunction{
		Value: func(args ...tengo.Object) (tengo.Object, error){
			gargs := make([]any, len(args))
			for i := range gargs {
				gargs[i] = tengo.ToInterface(args[i])
			}
			fmt.Fprint(&retBuf, gargs...)
			return nil, nil
			//return tengo.FromInterface([]byte(str))
		},
	})
	script.Add("__pp_println__", &tengo.UserFunction{
		Value: func(args ...tengo.Object) (tengo.Object, error){
			gargs := make([]any, len(args))
			for i := range gargs {
				gargs[i] = tengo.ToInterface(args[i])
			}
			fmt.Fprintln(&retBuf, gargs...)
			return nil, nil
			//return tengo.FromInterface([]byte(str))
		},
	})
	script.Add("__pp_write_raw__", &tengo.UserFunction{
		Value: func(args ...tengo.Object) (tengo.Object, error){
			bt := make([][]byte, len(args))
			for i, o := range args {
				bts, ok := tengo.ToByteSlice(o)
				if !ok {
					return nil, tengo.ErrInvalidArgumentType{
						Name: strconv.Itoa(i),
						Expected: "string/bytes",
						Found: o.TypeName(),
					}
				}
				bt[i] = bts
			}
			for _, b := range bt {
				retBuf.Write(b)
			}
			return nil, nil
		},
	})
	script.Add("__separate__", &tengo.UserFunction{
		Value: func(args ...tengo.Object) (tengo.Object, error){
			return nil, nil
			/*if len(args) < 1 {
				return nil, tengo.ErrWrongNumArguments
			}
			bts, ok := tengo.ToByteSlice(args[0])
			if !ok {
				return nil, tengo.ErrInvalidArgumentType{
				}
			}*/
			//rets = append(rets, retBuf.Bytes())
			//retBuf.Reset()
			return nil, nil
		},
	})

	compiled, err := script.Compile()
	if err != nil {
		return nil, nil, err
	}
	// To keep everything of the changes.
	compiled = compiled.Clone()

	_, err = script.RunContext(ctx)
	if err != nil {
		return nil, nil, err
	}

	return compiled, retBuf.Bytes(), nil
}