-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathserpent.go
221 lines (183 loc) · 5.8 KB
/
serpent.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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
// Package serpent provides functions for interacting with a Python interpreter.
package serpent
import (
"encoding/json"
"errors"
"fmt"
"io"
"os"
"runtime"
"sync"
"github.com/ebitengine/purego"
)
// Types used in the Python C API.
type pyObject uintptr
// Function prototypes for the Python C API.
var py_InitializeEx func(int)
var py_Finalize func()
var pyEval_GetBuiltins func() pyObject
var pyRun_String func(string, int, pyObject, pyObject) pyObject
var pyErr_Occurred func() bool
var pyErr_Print func()
var pyDict_New func() pyObject
var pyDict_GetItemString func(pyObject, string) pyObject
var pyDict_SetItemString func(pyObject, string, pyObject) int
var pyUnicode_AsUTF8 func(pyObject) string
var py_DecRef func(pyObject)
// Constants used in the Python C API.
const pyFileInput = 257
var (
// ErrAlreadyInitialized is returned when the Python interpreter is initialized more than once.
ErrAlreadyInitialized = errors.New("already initialized")
// ErrRunFailed is returned when the Python program fails to run.
ErrRunFailed = errors.New("run failed")
// ErrNoResult is returned when the result variable is not found in the Python program.
ErrNoResult = errors.New("no result")
)
// python is a handle to the Python shared library.
var python uintptr
// Init initializes the Python interpreter, loading the Python shared library from the supplied path. This
// must be called before any other functions in this package.
func Init(libraryPath string) error {
if python != 0 {
return ErrAlreadyInitialized
}
lib, err := purego.Dlopen(libraryPath, purego.RTLD_NOW|purego.RTLD_GLOBAL)
if err != nil {
return fmt.Errorf("dlopen: %v", err)
}
python = lib
purego.RegisterLibFunc(&py_InitializeEx, python, "Py_InitializeEx")
purego.RegisterLibFunc(&py_Finalize, python, "Py_Finalize")
purego.RegisterLibFunc(&pyEval_GetBuiltins, python, "PyEval_GetBuiltins")
purego.RegisterLibFunc(&pyErr_Occurred, python, "PyErr_Occurred")
purego.RegisterLibFunc(&pyErr_Print, python, "PyErr_Print")
purego.RegisterLibFunc(&pyDict_New, python, "PyDict_New")
purego.RegisterLibFunc(&pyDict_GetItemString, python, "PyDict_GetItemString")
purego.RegisterLibFunc(&pyDict_SetItemString, python, "PyDict_SetItemString")
purego.RegisterLibFunc(&pyUnicode_AsUTF8, python, "PyUnicode_AsUTF8")
purego.RegisterLibFunc(&py_DecRef, python, "Py_DecRef")
purego.RegisterLibFunc(&pyRun_String, python, "PyRun_String")
go start()
return nil
}
// Run runs a [Program] with the supplied argument and returns the result. The Python code must assign
// a result variable in the main program.
//
// Example Python program:
//
// result = input + 1
func Run[TInput, TResult any](program Program[TInput, TResult], arg TInput) (TResult, error) {
checkInit()
input, err := json.Marshal(arg)
if err != nil {
return *new(TResult), fmt.Errorf("marshal input: %w", err)
}
code := generateCode(string(program), input)
result, err := run(code)
if err != nil {
return *new(TResult), err
}
var value TResult
if err := json.Unmarshal([]byte(result), &value); err != nil {
return *new(TResult), fmt.Errorf("unmarshal result: %w", err)
}
return value, nil
}
// RunFd runs a [Program] with the supplied argument with the Python program writing to the supplied writer.
// The writer is made available as a file descriptor (fd) in the Python program. The Python program must close
// the file descriptor when it is done writing.
//
// Example Python program:
//
// import os
// os.write(fd, b'OK')
// os.close(fd)
func RunWrite[TInput any](w io.Writer, program Program[TInput, Writer], arg TInput) error {
checkInit()
pr, pw, err := os.Pipe()
if err != nil {
return fmt.Errorf("pipe: %w", err)
}
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
defer pr.Close()
io.Copy(w, pr)
}()
input, err := json.Marshal(struct {
Input TInput
Fd uintptr
}{arg, pw.Fd()})
if err != nil {
return fmt.Errorf("marshal input: %w", err)
}
code := generateWriterCode(string(program), input)
_, err = run(code)
if !errors.Is(err, ErrNoResult) {
return err
}
if err := pw.Close(); err != nil {
return fmt.Errorf("close writer: %w", err)
}
wg.Wait()
return nil
}
// runContext identifies the context of a Python run.
type runContext struct {
code string
cond *sync.Cond
done bool
value string
err error
}
// run runs a Python program and returns the result.
func run(code string) (string, error) {
var mu sync.Mutex
cond := sync.NewCond(&mu)
cond.L.Lock()
defer cond.L.Unlock()
ctx := &runContext{code: code, cond: cond}
runCh <- ctx
for !ctx.done {
cond.Wait()
}
return ctx.value, ctx.err
}
// runCh is a channel for sending Python code to the Python interpreter.
var runCh = make(chan *runContext, 1)
// start runs a loop waiting for instructions.
func start() {
runtime.LockOSThread()
defer runtime.UnlockOSThread()
py_InitializeEx(0)
defer py_Finalize()
for run := range runCh {
run.cond.L.Lock()
globals := pyDict_New()
pyDict_SetItemString(globals, "__builtins__", pyEval_GetBuiltins())
pyRun_String(run.code, pyFileInput, globals, globals)
if pyErr_Occurred() {
pyErr_Print()
run.err = ErrRunFailed
} else if item := pyDict_GetItemString(globals, "_result"); item != 0 {
run.value = pyUnicode_AsUTF8(item)
} else {
run.err = ErrNoResult
}
// This is a good candidate for sending the result on a channel, but doing so conflicts with Python's GIL.
// To work around that we set the result on the context and signal that the run is complete. The calling
// Run function waits for changes on the done state to know when the result is ready.
run.done = true
run.cond.Signal()
run.cond.L.Unlock()
py_DecRef(globals)
}
}
// checkInit checks if the Python interpreter has been initialized. It panics if it has not.
func checkInit() {
if python == 0 {
panic("serpent: Init must be called before Run")
}
}