146 lines
3.4 KiB
Go
146 lines
3.4 KiB
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) {
|
|
module := python3.PyImport_ImportModule("guessit") //ret val: new ref
|
|
if !(module != nil && python3.PyErr_Occurred() == nil) {
|
|
return nil, fmt.Errorf("failed to import 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")
|
|
}
|
|
|
|
fn := python3.PyDict_GetItemString(dict, "guessit") //retval: Borrowed
|
|
if !(fn != nil && python3.PyCallable_Check(fn)) {
|
|
return nil, fmt.Errorf("could not find function 'guessit'")
|
|
}
|
|
|
|
return &Guessit{
|
|
fn: fn,
|
|
}, 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()
|
|
|
|
testdataPy := g.fn.CallFunctionObjArgs(item, opts) //retval: New reference
|
|
if !(testdataPy != nil && python3.PyErr_Occurred() == nil) {
|
|
python3.PyErr_Print()
|
|
return
|
|
}
|
|
|
|
defer testdataPy.DecRef()
|
|
|
|
size := python3.PyDict_Size(testdataPy)
|
|
keys := python3.PyDict_Keys(testdataPy)
|
|
defer keys.DecRef()
|
|
|
|
vals := python3.PyDict_Values(testdataPy)
|
|
defer vals.DecRef()
|
|
|
|
tmpmap := make(map[string]interface{})
|
|
for i := 0; i < size; i++ {
|
|
kitem := python3.PyList_GetItem(keys, i)
|
|
key := python3.PyUnicode_AsUTF8(kitem)
|
|
kitem = nil
|
|
|
|
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)
|
|
tmp = append(tmp, v)
|
|
}
|
|
tmpmap[key] = tmp
|
|
default:
|
|
}
|
|
}
|
|
|
|
if _, ok := tmpmap["version"]; !ok {
|
|
tmpmap["version"] = -1
|
|
}
|
|
|
|
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
|
|
}
|