guessit-go/guessit.go

package guessit

// #cgo pkg-config: python-3.8-embed
// #include <Python.h>
// import "C"
import (
	"fmt"
	"os"
	"strings"

	python3 "git.stein-ivar.net/steino/cpy3"

	"github.com/mitchellh/mapstructure"
)

type Match struct {
	// Main
	Type             string `mapstructure:"type,omitempty"`
	Title            string `mapstructure:"title,omitempty"`
	ReleaseGroup     string `mapstructure:"release_group,omitempty"`
	StreamingService string `mapstructure:"streaming_service,omitempty"`

	// Episode
	Season  int64 `mapstructure:"season,omitempty"`
	Episode int64 `mapstructure:"episode,omitempty"`
	Version int64 `mapstructure:"version,omitempty"`

	// Video
	ScreenSize string `mapstructure:"screen_size,omitempty"`
	Container  string `mapstructure:"container,omitempty"`

	// Audio
	AudioChannels string   `mapstructure:"audio_channels,omitempty"`
	AudioCodec    []string `mapstructure:"audio_codec,omitempty"`

	// Other
	Other []string `mapstructure:"other,omitempty"`
}

type Guessit struct {
	fn *python3.PyObject
}

func init() {
	python3.Py_Initialize()
	if !python3.Py_IsInitialized() {
		fmt.Println("Error initializing the python interpreter")
		os.Exit(1)
	}
}

func New() (*Guessit, error) {
	oImport := python3.PyImport_ImportModule("guessit") //ret val: new ref
	if !(oImport != nil && python3.PyErr_Occurred() == nil) {
		return nil, fmt.Errorf("failed to import module 'guessit'")
	}

	module := python3.PyImport_AddModule("guessit") //ret val: borrowed ref (from oImport)
	if !(module != nil && python3.PyErr_Occurred() == nil) {
		return nil, fmt.Errorf("failed to add module 'guessit'")
	}
	defer module.DecRef()

	dict := python3.PyModule_GetDict(module) //ret val: Borrowed
	if !(dict != nil && python3.PyErr_Occurred() == nil) {
		return nil, fmt.Errorf("could not get dict for module")
	}

	guessitfn := python3.PyDict_GetItemString(dict, "guessit") //retval: Borrowed
	if !(guessitfn != nil && python3.PyCallable_Check(guessitfn)) {
		return nil, fmt.Errorf("could not find function 'guessit'")
	}
	return &Guessit{
		fn: guessitfn,
	}, nil
}

func (g Guessit) Guessit(s string, options ...string) (out Match, err error) {
	if len(s) == 0 {
		err = fmt.Errorf("input string is empty")
		return
	}

	item := python3.PyUnicode_FromString(s)
	defer item.DecRef()

	opts := python3.PyUnicode_FromString(strings.Join(options[:], " "))
	defer opts.DecRef()

	args := python3.PyTuple_New(2) //retval: New reference
	if args == nil {
		return
	}
	defer args.DecRef()

	ret := python3.PyTuple_SetItem(args, 0, item) //steals ref to pylist
	if ret != 0 {
		if python3.PyErr_Occurred() != nil {
			python3.PyErr_Print()
		}
		return
	}
	ret = python3.PyTuple_SetItem(args, 1, opts) //steals ref to pylist
	if ret != 0 {
		if python3.PyErr_Occurred() != nil {
			python3.PyErr_Print()
		}
		return
	}

	testdataPy := g.fn.CallObject(args) //retval: New reference
	if !(testdataPy != nil && python3.PyErr_Occurred() == nil) {
		python3.PyErr_Print()
		return
	}

	size := python3.PyDict_Size(testdataPy)
	keys := python3.PyDict_Keys(testdataPy)
	defer keys.DecRef()
	vals := python3.PyDict_Values(testdataPy)
	defer vals.DecRef()

	testdataPy = nil

	tmpmap := make(map[string]interface{})
	for i := 0; i < size; i++ {
		key := python3.PyUnicode_AsUTF8(python3.PyList_GetItem(keys, i))
		val := python3.PyList_GetItem(vals, i)
		switch {
		case python3.PyLong_Check(val):
			tmpmap[key] = python3.PyLong_AsLongLong(val)
		case python3.PyUnicode_Check(val):
			tmpmap[key] = python3.PyUnicode_AsUTF8(val)
		case python3.PyList_Check(val):
			var tmp []string
			for i := 0; i < python3.PyList_Size(val); i++ {
				item := python3.PyList_GetItem(val, i)
				v := python3.PyUnicode_AsUTF8(item)
				item.DecRef()
				tmp = append(tmp, v)
			}
			tmpmap[key] = tmp
		default:
		}
	}

	config := &mapstructure.DecoderConfig{
		WeaklyTypedInput: true,
		Result:           &out,
	}

	decoder, err := mapstructure.NewDecoder(config)
	if err != nil {
		return
	}

	err = decoder.Decode(tmpmap)
	if err != nil {
		return
	}

	return
}