Add PkgImporter to vm.Options for handling import statement

vm/vm.go

@@ -12,7 +12,8 @@ import (
 
 // Options provides options to run VM with
 type Options struct {
-	Debug bool // run in Debug mode
+	Debug       bool // run in Debug mode
+	PkgImporter Importer
 }
 
 type (

vm/vmExpr.go

@@ -522,7 +522,7 @@ func (runInfo *runInfoStruct) invokeExpr() {
 		name := runInfo.rv.String()
 		runInfo.rv = nilValue
 
-		methods, ok := env.Packages[name]
+		methods, ok := runInfo.options.Import(name)
 		if !ok {
 			runInfo.err = newStringError(expr, "package not found: "+name)
 			return
@@ -537,7 +537,7 @@ func (runInfo *runInfoStruct) invokeExpr() {
 			}
 		}
 
-		types, ok := env.PackageTypes[name]
+		types, ok := runInfo.options.ImportTypes(name)
 		if ok {
 			for typeName, typeValue := range types {
 				err = pack.DefineReflectType(typeName, typeValue)

vm/vmImport.go

@@ -0,0 +1,193 @@
+package vm
+
+import (
+	"reflect"
+	"sync"
+
+	"github.com/mattn/anko/env"
+)
+
+// Importer is an interface to handle `import` statement
+type Importer interface {
+	// Import package with designated `name`
+	Import(name string) (map[string]reflect.Value, bool)
+
+	// ImportTypes imports type definition for designated `name`
+	ImportTypes(name string) (map[string]reflect.Type, bool)
+
+	// Each iterate over available packages
+	Each(func(string, map[string]reflect.Value))
+
+	// EachTypes iterate over available package types
+	EachTypes(func(string, map[string]reflect.Type))
+
+	// AppendMap merge packages and package types into this Importer.
+	// If `pkgs` is not specified, ALL available packages and types will be merged.
+	// Otherwise, only designated `pkgs` are merged into this importer.
+	AppendMap(pkg map[string]map[string]reflect.Value, pkgTypes map[string]map[string]reflect.Type, pkgs ...string) Importer
+
+	// Append merge packages and package types into this Importer.
+	// If `pkgs` is not specified, ALL available packages and types will be merged.
+	// Otherwise, only designated `pkgs` are merged into this importer.
+	Append(i Importer, pkgs ...string) Importer
+
+	// Remove package and types specified in `pkgs`
+	Remove(pkgs ...string) Importer
+}
+
+type pkgImporter struct {
+	sync.RWMutex
+	pack      map[string]map[string]reflect.Value
+	packTypes map[string]map[string]reflect.Type
+}
+
+// NewPackagesImporter return import handler from map of packages and types
+func NewPackagesImporter(packages map[string]map[string]reflect.Value, packageTypes map[string]map[string]reflect.Type) Importer {
+	imp := pkgImporter{}
+	imp.copyMap(packages, packageTypes)
+	return &imp
+}
+
+// NewStdPackagesImporter create package importer for packages defined in `anko/packages`.
+// If `pkgs` is not set, ALL packages will be available for import.
+// Otherwise, if `pkgs` are specified, only designated packages will be available for import.
+func NewStdPackagesImporter(pkgs ...string) Importer {
+	imp := &pkgImporter{}
+	return imp.AppendMap(env.Packages, env.PackageTypes, pkgs...)
+
+}
+
+// NewPackagesWithStdImporter combines NewPackagesImporter and NewStdPackagesImporter.
+func NewPackagesWithStdImporter(packages map[string]map[string]reflect.Value, packageTypes map[string]map[string]reflect.Type, pkgs ...string) Importer {
+	imp := NewPackagesImporter(packages, packageTypes)
+	return imp.AppendMap(env.Packages, env.PackageTypes, pkgs...)
+}
+
+func (i *pkgImporter) Import(name string) (map[string]reflect.Value, bool) {
+	i.RLock()
+	defer i.RUnlock()
+
+	if i.pack == nil {
+		return nil, false
+	}
+	v, ok := i.pack[name]
+	return v, ok
+}
+func (i *pkgImporter) ImportTypes(name string) (map[string]reflect.Type, bool) {
+	i.RLock()
+	defer i.RUnlock()
+
+	if i.packTypes == nil {
+		return nil, false
+	}
+	t, ok := i.packTypes[name]
+	return t, ok
+}
+func (i *pkgImporter) Each(fn func(string, map[string]reflect.Value)) {
+	i.RLock()
+	for name, pkg := range i.pack {
+		fn(name, pkg)
+	}
+	i.RUnlock()
+}
+func (i *pkgImporter) EachTypes(fn func(string, map[string]reflect.Type)) {
+	i.RLock()
+	for name, typ := range i.packTypes {
+		fn(name, typ)
+	}
+	i.RUnlock()
+}
+func (i *pkgImporter) AppendMap(pkg map[string]map[string]reflect.Value, pkgTypes map[string]map[string]reflect.Type, pkgs ...string) Importer {
+	i.Lock()
+	i.makeMap()
+	for n, p := range pkg {
+		if i.contains(pkgs, n) {
+			i.pack[n] = p
+		}
+	}
+	for n, t := range pkgTypes {
+		if i.contains(pkgs, n) {
+			i.packTypes[n] = t
+		}
+	}
+	i.Unlock()
+
+	return i
+}
+
+func (i *pkgImporter) Append(other Importer, pkgs ...string) Importer {
+	i.Lock()
+	i.makeMap()
+	other.Each(func(name string, pkg map[string]reflect.Value) {
+		if i.contains(pkgs, name) {
+			i.pack[name] = pkg
+		}
+	})
+	other.EachTypes(func(name string, typ map[string]reflect.Type) {
+		if i.contains(pkgs, name) {
+			i.packTypes[name] = typ
+		}
+	})
+	i.Unlock()
+
+	return i
+}
+func (i *pkgImporter) Remove(pkgs ...string) Importer {
+	i.Lock()
+	for _, name := range pkgs {
+		delete(i.pack, name)
+		delete(i.packTypes, name)
+	}
+	i.Unlock()
+
+	return i
+}
+
+func (i *pkgImporter) contains(pkgs []string, name string) bool {
+	if len(pkgs) == 0 {
+		return true
+	}
+	for _, item := range pkgs {
+		if item == name {
+			return true
+		}
+	}
+	return false
+}
+
+func (i *pkgImporter) copyMap(pkg map[string]map[string]reflect.Value, pkgTypes map[string]map[string]reflect.Type) {
+	i.makeMap()
+	for n, p := range pkg {
+		i.pack[n] = p
+	}
+	for n, t := range pkgTypes {
+		i.packTypes[n] = t
+	}
+}
+
+func (i *pkgImporter) makeMap() {
+	if i.pack == nil {
+		i.pack = make(map[string]map[string]reflect.Value)
+	}
+	if i.packTypes == nil {
+		i.packTypes = make(map[string]map[string]reflect.Type)
+	}
+}
+
+// Import get packages from PkgImporter OR env.Packages.
+func (o *Options) Import(name string) (map[string]reflect.Value, bool) {
+	if o.PkgImporter != nil {
+		return o.PkgImporter.Import(name)
+	}
+	v, ok := env.Packages[name]
+	return v, ok
+}
+
+// ImportTypes get types from PkgImporter OR env.PackageTypes
+func (o *Options) ImportTypes(name string) (map[string]reflect.Type, bool) {
+	if o.PkgImporter != nil {
+		return o.PkgImporter.ImportTypes(name)
+	}
+	t, ok := env.PackageTypes[name]
+	return t, ok
+}

vm/vmImport_test.go

@@ -0,0 +1,76 @@
+package vm
+
+import (
+	"errors"
+	"fmt"
+	"reflect"
+	"testing"
+
+	_ "github.com/mattn/anko/packages"
+)
+
+func TestImportPackages(t *testing.T) {
+	type Descriptor struct {
+		ID   int
+		Name string
+		Desc string
+	}
+	var hello = func() {
+		fmt.Println("Msg: Hello, from 'custom' Package")
+	}
+	var sum = func(args ...float64) float64 {
+		s := 0.0
+		for _, v := range args {
+			s += v
+		}
+		return s
+	}
+
+	var pkg = map[string]map[string]reflect.Value{
+		"custom": {
+			"hello":  reflect.ValueOf(hello),
+			"sum":    reflect.ValueOf(sum),
+			"printf": reflect.ValueOf(fmt.Printf),
+		},
+	}
+	var typ = map[string]map[string]reflect.Type{
+		"custom": {
+			"Descriptor": reflect.TypeOf(Descriptor{}),
+		},
+	}
+
+	// Tests inputs
+	tNoPkg := Test{Script: `a = 1; b = 2; c = a+b;`, RunOutput: int64(3)}
+	tStdFmt := Test{Script: `var fmt = import("fmt"); fmt.Println("Msg: Hello from 'stdlib.fmt'"); x=0`, RunOutput: int64(0)}
+	tCust := Test{Script: `var custom = import("custom"); custom.hello(); s = custom.sum(1, 2, 3, 11); custom.printf("custom: sum = %f\n", s); out=s`, RunOutput: 17.0}
+	tCustFmt := Test{Script: `var custom, fmt = import("custom"), import("fmt"); custom.hello(); s = custom.sum(1, 2, 3, 11); fmt.Printf("stdlib.fmt: sum = %f\n", s); out=s`, RunOutput: 17.0}
+	tCustTimeErr := Test{Script: `var custom, time = import("custom"), import("time"); custom.hello(); now = time.Now(); custom.printf("Now %v\n", now); out=0`,
+		RunError: errors.New("package not found: time")}
+
+	// --- importer not specified (backward compatibility) ---
+	// should not error
+	runTests(t, []Test{tNoPkg, tStdFmt}, nil, nil)
+
+	// option without specifying package importer
+	runTests(t, []Test{tNoPkg, tStdFmt}, nil, &Options{Debug: true})
+
+	// --- explicitly specify importer --- //
+	// import only custom package
+	optCustom := Options{
+		PkgImporter: NewPackagesImporter(pkg, typ),
+	}
+	runTests(t, []Test{tNoPkg, tCust}, nil, &optCustom)
+
+	// import custom package and ALL stdlib
+	optCustomStdAll := Options{
+		PkgImporter: NewPackagesWithStdImporter(pkg, typ),
+	}
+	runTests(t, []Test{tNoPkg, tStdFmt, tCust, tCustFmt}, nil, &optCustomStdAll)
+
+	// import custom package and `fmt` from stdlib
+	optCustomStdFmt := Options{
+		PkgImporter: NewPackagesWithStdImporter(pkg, typ, "fmt"),
+	}
+	runTests(t, []Test{tNoPkg, tStdFmt, tCust, tCustFmt, tCustTimeErr}, nil, &optCustomStdFmt)
+
+}