1 files changed, 63 insertions, 7 deletions

diff --git a/vendor/github.com/d5/tengo/v2/compiler.go b/vendor/github.com/d5/tengo/v2/compiler.go
index 53cc7d38..e4e04303 100644
--- a/vendor/github.com/d5/tengo/v2/compiler.go
+++ b/vendor/github.com/d5/tengo/v2/compiler.go
@@ -1,9 +1,11 @@
 package tengo
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"reflect"
 	"strings"
@@ -45,11 +47,12 @@ type Compiler struct {
 	parent          *Compiler
 	modulePath      string
 	importDir       string
+	importFileExt   []string
 	constants       []Object
 	symbolTable     *SymbolTable
 	scopes          []compilationScope
 	scopeIndex      int
-	modules         *ModuleMap
+	modules         ModuleGetter
 	compiledModules map[string]*CompiledFunction
 	allowFileImport bool
 	loops           []*loop
@@ -63,7 +66,7 @@ func NewCompiler(
 	file *parser.SourceFile,
 	symbolTable *SymbolTable,
 	constants []Object,
-	modules *ModuleMap,
+	modules ModuleGetter,
 	trace io.Writer,
 ) *Compiler {
 	mainScope := compilationScope{
@@ -96,6 +99,7 @@ func NewCompiler(
 		trace:           trace,
 		modules:         modules,
 		compiledModules: make(map[string]*CompiledFunction),
+		importFileExt:   []string{SourceFileExtDefault},
 	}
 }
 
@@ -538,12 +542,8 @@ func (c *Compiler) Compile(node parser.Node) error {
 			}
 		} else if c.allowFileImport {
 			moduleName := node.ModuleName
-			if !strings.HasSuffix(moduleName, ".tengo") {
-				moduleName += ".tengo"
-			}
 
-			modulePath, err := filepath.Abs(
-				filepath.Join(c.importDir, moduleName))
+			modulePath, err := c.getPathModule(moduleName)
 			if err != nil {
 				return c.errorf(node, "module file path error: %s",
 					err.Error())
@@ -640,6 +640,39 @@ func (c *Compiler) SetImportDir(dir string) {
 	c.importDir = dir
 }
 
+// SetImportFileExt sets the extension name of the source file for loading
+// local module files.
+//
+// Use this method if you want other source file extension than ".tengo".
+//
+//     // this will search for *.tengo, *.foo, *.bar
+//     err := c.SetImportFileExt(".tengo", ".foo", ".bar")
+//
+// This function requires at least one argument, since it will replace the
+// current list of extension name.
+func (c *Compiler) SetImportFileExt(exts ...string) error {
+	if len(exts) == 0 {
+		return fmt.Errorf("missing arg: at least one argument is required")
+	}
+
+	for _, ext := range exts {
+		if ext != filepath.Ext(ext) || ext == "" {
+			return fmt.Errorf("invalid file extension: %s", ext)
+		}
+	}
+
+	c.importFileExt = exts // Replace the hole current extension list
+
+	return nil
+}
+
+// GetImportFileExt returns the current list of extension name.
+// Thease are the complementary suffix of the source file to search and load
+// local module files.
+func (c *Compiler) GetImportFileExt() []string {
+	return c.importFileExt
+}
+
 func (c *Compiler) compileAssign(
 	node parser.Node,
 	lhs, rhs []parser.Expr,
@@ -1098,6 +1131,7 @@ func (c *Compiler) fork(
 	child.parent = c              // parent to set to current compiler
 	child.allowFileImport = c.allowFileImport
 	child.importDir = c.importDir
+	child.importFileExt = c.importFileExt
 	if isFile && c.importDir != "" {
 		child.importDir = filepath.Dir(modulePath)
 	}
@@ -1287,6 +1321,28 @@ func (c *Compiler) printTrace(a ...interface{}) {
 	_, _ = fmt.Fprintln(c.trace, a...)
 }
 
+func (c *Compiler) getPathModule(moduleName string) (pathFile string, err error) {
+	for _, ext := range c.importFileExt {
+		nameFile := moduleName
+
+		if !strings.HasSuffix(nameFile, ext) {
+			nameFile += ext
+		}
+
+		pathFile, err = filepath.Abs(filepath.Join(c.importDir, nameFile))
+		if err != nil {
+			continue
+		}
+
+		// Check if file exists
+		if _, err := os.Stat(pathFile); !errors.Is(err, os.ErrNotExist) {
+			return pathFile, nil
+		}
+	}
+
+	return "", fmt.Errorf("module '%s' not found at: %s", moduleName, pathFile)
+}
+
 func resolveAssignLHS(
 	expr parser.Expr,
 ) (name string, selectors []parser.Expr) {