-
Notifications
You must be signed in to change notification settings - Fork 1.1k
/
helpers.go
179 lines (161 loc) · 4.46 KB
/
helpers.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
package python
/*
#include <stdlib.h>
typedef struct _pygilstate {} PyGILState_STATE;
PyGILState_STATE gilState;
*/
import "C"
import (
"errors"
"strings"
"unsafe"
)
const (
pythonPathKey = "PYTHONPATH"
)
// SetPythonPath is a helper for setting PYTHONPATH.
func SetPythonPath(p []string) {
mergedPaths := strings.Join(p, ":")
path := C.CString(mergedPaths)
defer C.free(unsafe.Pointer(path))
key := C.CString(pythonPathKey)
defer C.free(unsafe.Pointer(key))
C.setenv(key, path, 1)
}
// LoadModuleDict wraps PyModule_GetDict.
func LoadModuleDict(m string) (unsafe.Pointer, error) {
mod := C.CString(m)
defer C.free(unsafe.Pointer(mod))
modName := PyUnicode_FromString(mod)
if modName == nil {
return nil, errors.New("PyUnicode_FromString failed")
}
modObject := PyImport_Import(modName)
if modObject == nil {
return nil, errors.New("PyImport_Import failed")
}
dict := PyModule_GetDict(modObject)
if dict == nil {
return nil, errors.New("PyModule_GetDict failed")
}
return unsafe.Pointer(dict), nil
}
// GetItem wraps PyDict_GetItemString
func GetItem(d unsafe.Pointer, k string) (unsafe.Pointer, error) {
key := C.CString(k)
defer C.free(unsafe.Pointer(key))
obj := ToPyObject(d)
item := PyDict_GetItemString(obj, key)
if item == nil {
return nil, errors.New("GetItem failed")
}
return unsafe.Pointer(item), nil
}
// PyRunSimpleString wraps PyRun_SimpleStringFlags
func PyRunSimpleString(s string) {
cstr := C.CString(s)
defer C.free(unsafe.Pointer(cstr))
PyRun_SimpleStringFlags(cstr, nil)
}
// PyTupleNew wraps PyTuple_New
func PyTupleNew(size int) (unsafe.Pointer, error) {
tup := PyTuple_New(C.long(size))
if tup == nil {
return nil, errors.New("PyTupleNew failed")
}
return unsafe.Pointer(tup), nil
}
// PyTupleSetItem wraps PyTuple_SetItem
func PyTupleSetItem(tup unsafe.Pointer, pos int, o interface{}) error {
switch o.(type) {
case string:
str := C.CString(o.(string))
defer C.free(unsafe.Pointer(str))
pystr := PyUnicode_FromString(str)
if pystr == nil {
return errors.New("PyUnicode_FromString failed")
}
ret := PyTuple_SetItem(ToPyObject(tup), C.long(pos), pystr)
if ret != 0 {
return errors.New("PyTuple_SetItem failed")
}
default:
// Assume this is a PyObject
obj := o.(unsafe.Pointer)
ret := PyTuple_SetItem(ToPyObject(tup), C.long(pos), ToPyObject(obj))
if ret != 0 {
return errors.New("PyTuple_SetItem failed")
}
}
return nil
}
// PyTupleGetItem wraps PyTuple_GetItem
func PyTupleGetItem(tup unsafe.Pointer, pos int) (unsafe.Pointer, error) {
item := PyTuple_GetItem(ToPyObject(tup), C.long(pos))
if item == nil {
return nil, errors.New("PyTupleGetItem failed")
}
return unsafe.Pointer(item), nil
}
// PyObjectCallObject wraps PyObject_CallObject
func PyObjectCallObject(o unsafe.Pointer, args unsafe.Pointer) (unsafe.Pointer, error) {
ret := PyObject_CallObject(ToPyObject(o), ToPyObject(args))
if ret == nil {
return nil, errors.New("PyObjectCallObject failed")
}
return unsafe.Pointer(ret), nil
}
// PyObjectGetAttr wraps PyObject_GetAttr
func PyObjectGetAttr(o unsafe.Pointer, attr interface{}) (unsafe.Pointer, error) {
switch attr.(type) {
case string:
str := C.CString(attr.(string))
defer C.free(unsafe.Pointer(str))
pystr := PyUnicode_FromString(str)
if pystr == nil {
return nil, errors.New("PyUnicode_FromString failed")
}
ret := PyObject_GetAttr(ToPyObject(o), pystr)
if ret == nil {
return nil, errors.New("PyObjectGetAttr failed")
}
return unsafe.Pointer(ret), nil
}
return nil, nil
}
// PyBytesFromString wraps PyBytesFromString
func PyBytesFromString(input []byte) (unsafe.Pointer, error) {
data := C.CBytes(input)
defer C.free(unsafe.Pointer(data))
ret := PyBytes_FromStringAndSize((*C.char)(data), C.long(len(input)))
if ret == nil {
return nil, errors.New("PyBytesFromString failed")
}
return unsafe.Pointer(ret), nil
}
// PyBytesAsString wraps PyBytes_AsString
func PyBytesAsString(o unsafe.Pointer, l int) ([]byte, error) {
obj := ToPyObject(o)
cstr := PyBytes_AsString(obj)
if cstr == nil {
return nil, errors.New("PyBytes_AsString as string failed")
}
str := C.GoBytes(unsafe.Pointer(cstr), C.int(l))
b := []byte(str)
return b, nil
}
// PyLongAsLong wraps PyLong_AsLong
func PyLongAsLong(o unsafe.Pointer) int {
l := PyLong_AsLong(ToPyObject(o))
return int(l)
}
func PyIncRef(o unsafe.Pointer) {
Py_IncRef(ToPyObject(o))
}
func PyDecRef(o unsafe.Pointer) {
Py_DecRef(ToPyObject(o))
}
func PyTupleClearFreeList() int {
sz := PyTuple_ClearFreeList()
return int(sz)
}