jy-data-extract/src/jy/util/script.go

// script
package util

import (
	"encoding/json"
	"fmt"

	ljson "github.com/yuin/gopher-json"
	"github.com/yuin/gopher-lua"
)

type LuaScript struct {
	Code, Name, Script string
	Doc                map[string]interface{} //经过前置过滤的源信息
	Block              []*Block               //分块信息
	KvMap              map[string][]map[string]interface{}
	Result             map[string][]map[string]interface{} //抽取结果信息
	L                  *lua.LState
}

//stype per：前置,core：抽取，back：后置
func (s *LuaScript) RunScript(stype string) map[string]interface{} {
	data := map[string]interface{}{}
	s.L = lua.NewState()
	s.L.PreloadModule("json", ljson.Loader)
	defer s.L.Close()
	if err := s.L.DoString(s.Script); err != nil {
		data["err"] = err.Error()
	} else {
		if stype == "pre" {
			tab := MapToLuaTable(s.L, s.Doc)
			if err := s.L.CallByParam(lua.P{
				Fn:      s.L.GetGlobal("main"),
				NRet:    1,
				Protect: true,
			}, tab); err != nil {
				data["err"] = err.Error()
			}
		} else if stype == "core" {
			tab := MapToLuaTable(s.L, s.Doc)
			block, _ := json.Marshal(s.Block)
			kvMap := MapToLuaTable2(s.L, s.KvMap)
			if err := s.L.CallByParam(lua.P{
				Fn:      s.L.GetGlobal("main"),
				NRet:    1,
				Protect: true,
			}, tab, lua.LString(block), kvMap); err != nil {
				data["err"] = err.Error()
			}
		} else if stype == "back" {
			result := MapToLuaTable2(s.L, s.Result)
			if err := s.L.CallByParam(lua.P{
				Fn:      s.L.GetGlobal("main"),
				NRet:    1,
				Protect: true,
			}, result); err != nil {
				data["err"] = err.Error()
			}
		}
		ret := s.L.Get(-1)
		s.L.Pop(1)
		if tmp, ok := ret.(*lua.LTable); ok {
			data = LuaTableToMap(tmp)
		} else {
			data["rep"] = ret
		}
	}
	return data
}
func Logic(str string, doc map[string]interface{}) bool {
	L := lua.NewState()
	L.PreloadModule("json", ljson.Loader)
	defer L.Close()
	b := false
	if err := L.DoString(str); err != nil {
		panic(err)
	} else {
		tab := MapToLuaTable(L, doc)
		if err := L.CallByParam(lua.P{
			Fn:      L.GetGlobal("logic"),
			NRet:    1,
			Protect: true,
		}, tab); err != nil {
			panic(err)
		}
		ret := L.Get(-1)
		L.Pop(1)
		if ret.String() == "true" {
			b = true
		}
	}
	return b
}

func MapToLuaTable(l *lua.LState, obj map[string]interface{}) *lua.LTable {
	tab := l.NewTable()
	for k, v := range obj {
		if val, ok := v.(string); ok {
			tab.RawSet(lua.LString(k), lua.LString(val))
		} else if val, ok := v.(int64); ok {
			tab.RawSet(lua.LString(k), lua.LNumber(val))
		} else if val, ok := v.(int32); ok {
			tab.RawSet(lua.LString(k), lua.LNumber(val))
		} else if val, ok := v.(float64); ok {
			tab.RawSet(lua.LString(k), lua.LNumber(val))
		} else if val, ok := v.(float32); ok {
			tab.RawSet(lua.LString(k), lua.LNumber(val))
		} else if val, ok := v.(bool); ok {
			tab.RawSet(lua.LString(k), lua.LBool(val))
		} else if val, ok := v.(map[string]interface{}); ok {
			tab.RawSet(lua.LString(k), MapToLuaTable(l, val))
		} else if val, ok := v.([]string); ok {
			tb := l.NewTable()
			for k, vv := range val {
				tb.Insert(k, lua.LString(vv))
			}
			tab.RawSet(lua.LString(k), tb)
		} else if val, ok := v.([]map[string]interface{}); ok {
			tab2 := l.NewTable()
			for i, v := range val {
				tab2.Insert(i, MapToLuaTable(l, v))
			}
			tab.RawSet(lua.LString(k), tab2)
		} else if val, ok := v.([]interface{}); ok {
			bs, _ := json.Marshal(val)
			tab.RawSet(lua.LString(k), lua.LString(string(bs)))
		}
	}
	return tab
}

func MapToLuaTable2(l *lua.LState, obj map[string][]map[string]interface{}) *lua.LTable {
	tab := l.NewTable()
	for k, ms := range obj {
		tab2 := l.NewTable()
		for i, v := range ms {
			tab2.Insert(i+1 /*加1防止顺序错乱*/, MapToLuaTable(l, v))
		}
		tab.RawSet(lua.LString(k), tab2)
	}
	return tab
}

func LuaTableToMap(param *lua.LTable) map[string]interface{} {
	tmp := map[string]interface{}{}
	param.ForEach(func(key, val lua.LValue) {
		kk := fmt.Sprint(key)
		if v, ok := val.(lua.LString); ok {
			tmp[kk] = string(v)
		} else if v, ok := val.(*lua.LTable); ok {
			t := []map[string]interface{}{}
			v.ForEach(func(k, inv lua.LValue) {
				if vv, ok := inv.(*lua.LTable); ok {
					t = append(t, LuaTableToMap(vv))
				}
			})
			tmp[kk] = t
		} else if v, ok := val.(*lua.LBool); ok {
			if v.String() == "true" {
				tmp[kk] = true
			} else {
				tmp[kk] = false
			}
		} else {
			tmp[kk] = v
		}
	})
	return tmp
}