add custom extension support for importing source file (#350)

* chore: add tests for custom extension

* feat: cusom source extension #286

* fix: path to test directory

* add getter + change setter name for file extension

* add tests of source file extension validation

* fix: add validation for file extension names

* fix: property importExt -> importFileExt

* fix: redundant check (no resetting)

* fix: failing test wich did not follow the new spec

* chore: add detailed description of the test

* chore: fix doc comment to be descriptive

* docs: add note about customizing the file extension

compiler.go

@@ -1,9 +1,11 @@
 package tengo
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"reflect"
 	"strings"
@@ -45,6 +47,7 @@ type Compiler struct {
 	parent          *Compiler
 	modulePath      string
 	importDir       string
+	importFileExt   []string
 	constants       []Object
 	symbolTable     *SymbolTable
 	scopes          []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) {

compiler_test.go

@@ -2,12 +2,15 @@ package tengo_test
 
 import (
 	"fmt"
+	"io/ioutil"
+	"path/filepath"
 	"strings"
 	"testing"
 
 	"github.com/d5/tengo/v2"
 	"github.com/d5/tengo/v2/parser"
 	"github.com/d5/tengo/v2/require"
+	"github.com/d5/tengo/v2/stdlib"
 )
 
 func TestCompiler_Compile(t *testing.T) {
@@ -1010,7 +1013,7 @@ r["x"] = {
 	expectCompileError(t, `
 (func() {
 	fn := fn()
-})()	
+})()
 `, "unresolved reference 'fn")
 }
 
@@ -1222,6 +1225,91 @@ func() {
 				tengo.MakeInstruction(parser.OpReturn, 0)))))
 }
 
+func TestCompiler_custom_extension(t *testing.T) {
+	pathFileSource := "./testdata/issue286/test.mshk"
+
+	modules := stdlib.GetModuleMap(stdlib.AllModuleNames()...)
+
+	src, err := ioutil.ReadFile(pathFileSource)
+	require.NoError(t, err)
+
+	// Escape shegang
+	if len(src) > 1 && string(src[:2]) == "#!" {
+		copy(src, "//")
+	}
+
+	fileSet := parser.NewFileSet()
+	srcFile := fileSet.AddFile(filepath.Base(pathFileSource), -1, len(src))
+
+	p := parser.NewParser(srcFile, src, nil)
+	file, err := p.ParseFile()
+	require.NoError(t, err)
+
+	c := tengo.NewCompiler(srcFile, nil, nil, modules, nil)
+	c.EnableFileImport(true)
+	c.SetImportDir(filepath.Dir(pathFileSource))
+
+	// Search for "*.tengo" and ".mshk"(custom extension)
+	c.SetImportFileExt(".tengo", ".mshk")
+
+	err = c.Compile(file)
+	require.NoError(t, err)
+}
+
+func TestCompilerNewCompiler_default_file_extension(t *testing.T) {
+	modules := stdlib.GetModuleMap(stdlib.AllModuleNames()...)
+	input := "{}"
+	fileSet := parser.NewFileSet()
+	file := fileSet.AddFile("test", -1, len(input))
+
+	c := tengo.NewCompiler(file, nil, nil, modules, nil)
+	c.EnableFileImport(true)
+
+	require.Equal(t, []string{".tengo"}, c.GetImportFileExt(),
+		"newly created compiler object must contain the default extension")
+}
+
+func TestCompilerSetImportExt_extension_name_validation(t *testing.T) {
+	c := new(tengo.Compiler) // Instantiate a new compiler object with no initialization
+
+	// Test of empty arg
+	err := c.SetImportFileExt()
+
+	require.Error(t, err, "empty arg should return an error")
+
+	// Test of various arg types
+	for _, test := range []struct {
+		extensions []string
+		expect     []string
+		requireErr bool
+		msgFail    string
+	}{
+		{[]string{".tengo"}, []string{".tengo"}, false,
+			"well-formed extension should not return an error"},
+		{[]string{""}, []string{".tengo"}, true,
+			"empty extension name should return an error"},
+		{[]string{"foo"}, []string{".tengo"}, true,
+			"name without dot prefix should return an error"},
+		{[]string{"foo.bar"}, []string{".tengo"}, true,
+			"malformed extension should return an error"},
+		{[]string{"foo."}, []string{".tengo"}, true,
+			"malformed extension should return an error"},
+		{[]string{".mshk"}, []string{".mshk"}, false,
+			"name with dot prefix should be added"},
+		{[]string{".foo", ".bar"}, []string{".foo", ".bar"}, false,
+			"it should replace instead of appending"},
+	} {
+		err := c.SetImportFileExt(test.extensions...)
+		if test.requireErr {
+			require.Error(t, err, test.msgFail)
+		}
+
+		expect := test.expect
+		actual := c.GetImportFileExt()
+		require.Equal(t, expect, actual, test.msgFail)
+	}
+}
+
 func concatInsts(instructions ...[]byte) []byte {
 	var concat []byte
 	for _, i := range instructions {

docs/tutorial.md

@@ -37,22 +37,22 @@ Here's a list of all available value types in Tengo.
 | map | value map with string keys _(mutable)_ | `map[string]interface{}` |
 | immutable map | [immutable](#immutable-values) map | - |
 | undefined | [undefined](#undefined-values) value | - |
-| function | [function](#function-values) value | - |  
+| function | [function](#function-values) value | - |
 | _user-defined_ | value of [user-defined types](https://github.com/d5/tengo/blob/master/docs/objects.md) | - |
 
 ### Error Values
 
 In Tengo, an error can be represented using "error" typed values. An error
 value is created using `error` expression, and, it must have an underlying
 value. The underlying value of an error value can be access using `.value`
-selector.  
+selector.
 
 ```golang
 err1 := error("oops")    // error with string value
 err2 := error(1+2+3)     // error with int value
 if is_error(err1) {      // 'is_error' builtin function
   err_val := err1.value  // get underlying value
-}  
+}
 ```
 
 ### Immutable Values
@@ -101,12 +101,12 @@ a.c[1] = 5     // illegal
 ### Undefined Values
 
 In Tengo, an "undefined" value can be used to represent an unexpected or
-non-existing value:  
+non-existing value:
 
 - A function that does not return a value explicitly considered to return
 `undefined` value.
 - Indexer or selector on composite value types may return `undefined` if the
-key or index does not exist.  
+key or index does not exist.
 - Type conversion builtin functions without a default value will return
 `undefined` if conversion fails.
 
@@ -142,8 +142,8 @@ m["b"]                                // == false
 m.c                                   // == "foo"
 m.x                                   // == undefined
 
-{a: [1,2,3], b: {c: "foo", d: "bar"}} // ok: map with an array element and a map element  
-```  
+{a: [1,2,3], b: {c: "foo", d: "bar"}} // ok: map with an array element and a map element
+```
 
 ### Function Values
 
@@ -233,7 +233,7 @@ a := "foo"      // define 'a' in global scope
 
 func() {        // function scope A
   b := 52       // define 'b' in function scope A
-  
+
   func() {      // function scope B
     c := 19.84  // define 'c' in function scope B
 
@@ -243,12 +243,12 @@ func() {        // function scope A
     b := true   // ok: define new 'b' in function scope B
                 //     (shadowing 'b' from function scope A)
   }
-  
+
   a = "bar"     // ok: assigne new value to 'a' from global scope
   b = 10        // ok: assigne new value to 'b'
   a := -100     // ok: define new 'a' in function scope A
                 //     (shadowing 'a' from global scope)
-  
+
   c = -9.1      // illegal: 'c' is not defined
   b := [1, 2]   // illegal: 'b' is already defined in the same scope
 }
@@ -470,15 +470,15 @@ for {
 
 "For-In" statement is new in Tengo. It's similar to Go's `for range` statement.
 "For-In" statement can iterate any iterable value types (array, map, bytes,
-string, undefined).  
+string, undefined).
 
 ```golang
 for v in [1, 2, 3] {          // array: element
   // 'v' is value
 }
 for i, v in [1, 2, 3] {       // array: index and element
   // 'i' is index
-  // 'v' is value  
+  // 'v' is value
 }
 for k, v in {k1: 1, k2: 2} {  // map: key and value
   // 'k' is key
@@ -508,6 +508,16 @@ export func(x) {
 }
 ```
 
+By default, `import` solves the missing extension name of a module file as
+"`.tengo`"[^note].
+Thus, `sum := import("./sum")` is equivalent to `sum := import("./sum.tengo")`.
+
+[^note]:
+    If using Tengo as a library in Go, the file extension name "`.tengo`" can
+    be customized. In that case, use the `SetImportFileExt` function of the
+    `Compiler` type.
+    See the [Go reference](https://pkg.go.dev/github.com/d5/tengo/v2) for details.
+
 In Tengo, modules are very similar to functions.
 
 - `import` expression loads the module code and execute it like a function.
@@ -517,9 +527,9 @@ In Tengo, modules are very similar to functions.
   return a value to the importing code.
   - `export`-ed values are always immutable.
   - If the module does not have any `export` statement, `import` expression
-  simply returns `undefined`. _(Just like the function that has no `return`.)_  
+  simply returns `undefined`. _(Just like the function that has no `return`.)_
   - Note that `export` statement is completely ignored and not evaluated if
-  the code is executed as a main module.  
+  the code is executed as a main module.
 
 Also, you can use `import` expression to load the
 [Standard Library](https://github.com/d5/tengo/blob/master/docs/stdlib.md) as

tengo.go

@@ -26,6 +26,9 @@ const (
 
 	// MaxFrames is the maximum number of function frames for a VM.
 	MaxFrames = 1024
+
+	// SourceFileExtDefault is the default extension for source files.
+	SourceFileExtDefault = ".tengo"
 )
 
 // CallableFunc is a function signature for the callable functions.

testdata/issue286/dos/cinco/cinco.mshk

@@ -0,0 +1,8 @@
+export {
+    fn: func(...args) {
+        text := import("text")
+        args = append(args, "cinco")
+
+        return text.join(args, " ")
+    }
+}

testdata/issue286/dos/dos.mshk

@@ -0,0 +1,7 @@
+export {
+    fn: func(a, b) {
+        tres := import("../tres")
+
+        return tres.fn(a, b, "dos")
+    }
+}

testdata/issue286/dos/quatro/quatro.mshk

@@ -0,0 +1,7 @@
+export {
+    fn: func(a, b, c, d) {
+        cinco := import("../cinco/cinco")
+
+        return cinco.fn(a, b, c, d, "quatro")
+    }
+}