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

// script
package util

import (
	"encoding/json"
	"fmt"
	"log"
	qu "qfw/util"
	"regexp"
	"strings"

	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
}

func (s *LuaScript) ClearRunScript() map[string]interface{} {
	data := map[string]interface{}{}
	qu.Try(func() {
		s.L = lua.NewState() //创建lua虚拟机
		s.L.PreloadModule("json", ljson.Loader)
		defer s.L.Close()
		if err := s.L.DoString(s.Script); err != nil {
			data["err"] = err.Error()
		} else {
			//doc:每条结果数据
			//KvMap:合并数据的list
			tab := MapToLuaTable(s.L, s.Doc)
			kvMap := MapToLuaTable2(s.L, s.KvMap)
			if err := s.L.CallByParam(lua.P{
				Fn:      s.L.GetGlobal("main"),
				NRet:    1,
				Protect: true,
			}, lua.LString(s.Code), tab, kvMap); 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
			}
		}
	}, func(err interface{}) {
		log.Println("lua err:", data["err"])
	})
	return data
}

//stype per：前置,core：抽取，back：后置
func (s *LuaScript) RunScript(stype string) map[string]interface{} {
	data := map[string]interface{}{}
	qu.Try(func() {
		s.L = lua.NewState()
		//s.L.DoFile("lua/comm.lua")
		s.L.PreloadModule("json", ljson.Loader)
		s.L.SetGlobal("findOneText", s.L.NewFunction(func(S *lua.LState) int {
			nodetype := S.ToString(-3)
			gpath := S.ToString(-2)
			content := S.ToString(-1)
			ret := FindOneText(gpath, content, nodetype)
			S.Push(ret)
			return 1
		}))
		s.L.SetGlobal("findOneHtml", s.L.NewFunction(func(S *lua.LState) int {
			nodetype := S.ToString(-3)
			gpath := S.ToString(-2)
			content := S.ToString(-1)
			ret := FindOneHtml(gpath, content, nodetype)
			S.Push(ret)
			return 1
		}))
		s.L.SetGlobal("findListText", s.L.NewFunction(func(S *lua.LState) int {
			gpath := S.ToString(-2)
			content := S.ToString(-1)
			ret := s.L.NewTable()
			FindListText(gpath, content, ret)
			S.Push(ret)
			return 1
		}))
		s.L.SetGlobal("findListHtml", s.L.NewFunction(func(S *lua.LState) int {
			gpath := S.ToString(-2)
			content := S.ToString(-1)
			ret := s.L.NewTable()
			FindListHtml(gpath, content, ret)
			S.Push(ret)
			return 1
		}))
		s.L.SetGlobal("findMap", s.L.NewFunction(func(S *lua.LState) int {
			qmap := S.ToTable(-2)
			content := S.ToString(-1)
			ret := s.L.NewTable()
			FindMap(qmap, content, ret)
			S.Push(ret)
			return 1
		}))
		//支持正则
		s.L.SetGlobal("regexp", s.L.NewFunction(func(S *lua.LState) int {
			index := int(S.ToNumber(-1))
			regstr := S.ToString(-2)
			text := S.ToString(-3)
			reg := regexp.MustCompile(regstr)
			reps := reg.FindAllStringSubmatchIndex(text, -1)
			ret := s.L.NewTable()
			number := 0
			for _, v := range reps {
				number++
				ret.Insert(number, lua.LString(text[v[index]:v[index+1]]))
			}
			S.Push(ret)
			return 1
		}))
		//支持替换
		s.L.SetGlobal("replace", s.L.NewFunction(func(S *lua.LState) int {
			text := S.ToString(-3)
			old := S.ToString(-2)
			repl := S.ToString(-1)
			text = strings.Replace(text, old, repl, -1)
			S.Push(lua.LString(text))
			return 1
		}))
		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,
				}, lua.LString(s.Code), 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,
				}, lua.LString(s.Code), tab, lua.LString(block), kvMap); err != nil {
					data["err"] = err.Error()
					log.Println(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,
				}, lua.LString(s.Code), 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
			}
		}
	}, func(err interface{}) {
		log.Println("lua err:", err)
	})
	return data
}
func Logic(str string, doc map[string]interface{}) bool {
	b := false
	var errs error
	qu.Try(func() {
		L := lua.NewState()
		L.PreloadModule("json", ljson.Loader)
		defer L.Close()
		if err := L.DoString(str); err != nil {
			errs = err
		} else {
			tab := MapToLuaTable(L, doc)
			if err := L.CallByParam(lua.P{
				Fn:      L.GetGlobal("logic"),
				NRet:    1,
				Protect: true,
			}, tab); err != nil {
				errs = err
			}
			ret := L.Get(-1)
			L.Pop(1)
			if ret.String() == "true" {
				b = true
			}
		}
	}, func(err interface{}) {
		log.Println("lua err:", errs)
	})
	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.(int); 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]bool);ok{//把map的value值为true的存储为数组
			tb := l.NewTable()
			var i int
			for k2,_ := range val{
				tb.Insert(i, lua.LString(k2))
				i+=1
			}
			tab.RawSet(lua.LString(k), tb)
		} 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 {
			if kk == "blocktag"{//分块
				t := map[string]bool{}
				v.ForEach(func(value lua.LValue, value2 lua.LValue) {
					t[value2.String()]=true
				})
				tmp[kk] = t
				return
			}
			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 if kk == "comeintime" || kk == "publishtime" {
			tmp[kk] = qu.Int64All(val.String())
		} else if v, ok := val.(*lua.LNumber); ok {
			tmp[kk] = qu.Float64All(v.String())
		} else if v, ok := val.(lua.LNumber); ok {
			tmp[kk] = qu.Float64All(v.String())
		} else {
			tmp[kk] = val
		}
	})
	return tmp
}