File: pickle_machine.go

package info (click to toggle)
golang-github-hydrogen18-stalecucumber 0.0~git20161112.0.1e918be-1
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 240 kB
  • ctags: 350
  • sloc: python: 71; makefile: 2
file content (598 lines) | stat: -rw-r--r-- 15,556 bytes parent folder | download | duplicates (3)
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
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
/*
This package reads and writes pickled data. The format is the same
as the Python "pickle" module.

Protocols 0,1,2 are implemented. These are the versions written by the Python
2.x series. Python 3 defines newer protocol versions, but can write the older
protocol versions so they are readable by this package.

To read data, see stalecucumber.Unpickle.

To write data, see stalecucumber.NewPickler.

TLDR

Read a pickled string or unicode object
	pickle.dumps("foobar")
	---
	var somePickledData io.Reader
	mystring, err := stalecucumber.String(stalecucumber.Unpickle(somePickledData))

Read a pickled integer
	pickle.dumps(42)
	---
	var somePickledData io.Reader
	myint64, err := stalecucumber.Int(stalecucumber.Unpickle(somePickledData))

Read a pickled list of numbers into a structure
	pickle.dumps([8,8,2005])
	---
	var somePickledData io.Reader
	numbers := make([]int64,0)

	err := stalecucumber.UnpackInto(&numbers).From(stalecucumber.Unpickle(somePickledData))

Read a pickled dictionary into a structure
	pickle.dumps({"apple":1,"banana":2,"cat":"hello","Dog":42.0})
	---
	var somePickledData io.Reader
	mystruct := struct{
		Apple int
		Banana uint
		Cat string
		Dog float32}{}

	err := stalecucumber.UnpackInto(&mystruct).From(stalecucumber.Unpickle(somePickledData))

Pickle a struct

	buf := new(bytes.Buffer)
	mystruct := struct{
			Apple int
			Banana uint
			Cat string
			Dog float32}{}

	err := stalecucumber.NewPickler(buf).Pickle(mystruct)



Recursive objects

You can pickle recursive objects like so

	a = {}
	a["self"] = a
	pickle.dumps(a)

Python's pickler is intelligent enough not to emit an infinite data structure
when a recursive object is pickled.

I recommend against pickling recursive objects in the first place, but this
library handles unpickling them without a problem. The result of unpickling
the above is map[interface{}]interface{} with a key "a" that contains
a reference to itself.

Attempting to unpack the result of the above python code into a structure
with UnpackInto would either fail or recurse forever.

Protocol Performance

If the version of Python you are using supports protocol version 1 or 2,
you should always specify that protocol version. By default the "pickle"
and "cPickle" modules in Python write using protocol 0. Protocol 0
requires much more space to represent the same values and is much
slower to parse.

Unsupported Opcodes

The pickle format is incredibly flexible and as a result has some
features that are impractical or unimportant when implementing a reader in
another language.

Each set of opcodes is listed below by protocol version with the impact.

Protocol 0

	GLOBAL

This opcode is equivalent to calling "import foo; foo.bar" in python. It is
generated whenever an object instance, class definition, or method definition
is serialized. As long as the pickled data does not contain an instance
of a python class or a reference to a python callable this opcode is not
emitted by the "pickle" module.

A few examples of what will definitely cause this opcode to be emitted

	pickle.dumps(range) #Pickling the range function
	pickle.dumps(Exception()) #Pickling an instance of a python class

This opcode will be partially supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	REDUCE
	BUILD
	INST

These opcodes are used in recreating pickled python objects. That is currently
not supported by this package.

These opcodes will be supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	PERSID

This opcode is used to reference concrete definitions of objects between
a pickler and an unpickler by an ID number. The pickle protocol doesn't define
what a persistent ID means.

This opcode is unlikely to ever be supported by this package.

Protocol 1

	OBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.


	BINPERSID

This opcode is equivalent to PERSID in protocol 0 and won't be supported
for the same reason.

Protocol 2

	NEWOBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	EXT1
	EXT2
	EXT4

These opcodes allow using a registry
of popular objects that are pickled by name, typically classes.
It is envisioned that through a global negotiation and
registration process, third parties can set up a mapping between
ints and object names.

These opcodes are unlikely to ever be supported by this package.

*/
package stalecucumber

import "errors"
import "io"
import "bytes"
import "encoding/binary"
import "fmt"

var ErrOpcodeStopped = errors.New("STOP opcode found")
var ErrStackTooSmall = errors.New("Stack is too small to perform requested operation")
var ErrInputTruncated = errors.New("Input to the pickle machine was truncated")
var ErrOpcodeNotImplemented = errors.New("Input encountered opcode that is not implemented")
var ErrNoResult = errors.New("Input did not place a value onto the stack")
var ErrMarkNotFound = errors.New("Mark could not be found on the stack")

/*
Unpickle a value from a reader. This function takes a reader and
attempts to read a complete pickle program from it. This is normally
the output of the function "pickle.dump" from Python.

The returned type is interface{} because unpickling can generate any type. Use
a helper function to convert to another type without an additional type check.

This function returns an error if
the reader fails, the pickled data is invalid, or if the pickled data contains
an unsupported opcode. See unsupported opcodes in the documentation of
this package for more information.

Type Conversions

Types conversion Python types to Go types is performed as followed
	int -> int64
	string -> string
	unicode -> string
	float -> float64
	long -> big.Int from the "math/big" package
	lists -> []interface{}
	tuples -> []interface{}
	dict -> map[interface{}]interface{}

The following values are converted from Python to the Go types
	True & False -> bool
	None -> stalecucumber.PickleNone, sets pointers to nil

Helper Functions

The following helper functions were inspired by the github.com/garyburd/redigo
package. Each function takes the result of Unpickle as its arguments. If unpickle
fails it does nothing and returns that error. Otherwise it attempts to
convert to the appropriate type. If type conversion fails it returns an error

	String - string from Python string or unicode
	Int - int64 from Python int or long
	Bool - bool from Python True or False
	Big - *big.Int from Python long
	ListOrTuple - []interface{} from Python Tuple or List
	Float - float64 from Python float
	Dict - map[interface{}]interface{} from Python dictionary
	DictString -
		map[string]interface{} from Python dictionary.
		Keys must all be of type unicode or string.

Unpacking into structures

If the pickled object is a python dictionary that has only unicode and string
objects for keys, that object can be unpickled into a struct in Go by using
the "UnpackInto" function. The "From" receiver on the return value accepts
the result of "Unpickle" as its actual parameters.

The keys of the python dictionary are assigned to fields in a structure.
Structures may specify the tag "pickle" on fields. The value of this tag is taken
as the key name of the Python dictionary value to place in this field. If no
field has a matching "pickle" tag the fields are looked up by name. If
the first character of the key is not uppercase, it is uppercased. If a field
matching that name is found, the value in the python dictionary is unpacked
into the value of the field within the structure.

A list of python dictionaries can be unpickled into a slice of structures in
Go.

A homogeneous list of python values can be unpickled into a slice in
Go with the appropriate element type.

A nested python dictionary is unpickled into nested structures in Go. If a
field is of type map[interface{}]interface{} it is of course unpacked into that
as well.

By default UnpackInto skips any missing fields and fails if a field's
type is not compatible with the object's type.

This behavior can be changed by setting "AllowMissingFields" and
"AllowMismatchedFields" on the return value of UnpackInto before calling
From.

*/
func Unpickle(reader io.Reader) (interface{}, error) {
	var pm PickleMachine
	pm.buf = &bytes.Buffer{}
	pm.Reader = reader
	pm.lastMark = -1
	//Pre allocate a small stack
	pm.Stack = make([]interface{}, 0, 16)

	err := (&pm).execute()
	if err != ErrOpcodeStopped {
		return nil, pm.error(err)
	}

	if len(pm.Stack) == 0 {
		return nil, ErrNoResult
	}

	return pm.Stack[0], nil

}

var jumpList = buildEmptyJumpList()

func init() {
	populateJumpList(&jumpList)
}

/*
This type is returned whenever Unpickle encounters an error in pickled data.
*/
type PickleMachineError struct {
	Err       error
	StackSize int
	MemoSize  int
	Opcode    uint8
}

/*
This struct is current exposed but not useful. It is likely to be hidden
in the near future.
*/
type PickleMachine struct {
	Stack  []interface{}
	Memo   []interface{}
	Reader io.Reader

	currentOpcode uint8
	buf           *bytes.Buffer
	lastMark      int

	memoBuffer               [16]memoBufferElement
	memoBufferMaxDestination int64
	memoBufferIndex          int
}

type memoBufferElement struct {
	Destination int64
	V           interface{}
}

func (pme PickleMachineError) Error() string {
	return fmt.Sprintf("Pickle Machine failed on opcode:0x%x. Stack size:%d. Memo size:%d. Cause:%v", pme.Opcode, pme.StackSize, pme.MemoSize, pme.Err)
}

func (pm *PickleMachine) error(src error) error {
	return PickleMachineError{
		StackSize: len(pm.Stack),
		MemoSize:  len(pm.Memo),
		Err:       src,
		Opcode:    pm.currentOpcode,
	}
}

func (pm *PickleMachine) execute() error {
	for {
		err := binary.Read(pm.Reader, binary.BigEndian, &pm.currentOpcode)
		if err != nil {
			return err
		}

		err = jumpList[int(pm.currentOpcode)](pm)

		if err != nil {
			return err
		}
	}
}

func (pm *PickleMachine) flushMemoBuffer(vIndex int64, v interface{}) {
	//Extend the memo until it is large enough
	if pm.memoBufferMaxDestination >= int64(len(pm.Memo)) {
		replacement := make([]interface{}, pm.memoBufferMaxDestination<<1)
		copy(replacement, pm.Memo)
		pm.Memo = replacement
	}

	//If a value was passed into this function, write it into the memo
	//as well
	if vIndex != -1 {
		pm.Memo[vIndex] = v
	}

	//Write the contents of the buffer into the memo
	//in the same order as the puts were issued
	for i := 0; i != pm.memoBufferIndex; i++ {
		buffered := pm.memoBuffer[i]
		pm.Memo[buffered.Destination] = buffered.V
	}

	//Reset the buffer
	pm.memoBufferIndex = 0
	pm.memoBufferMaxDestination = 0

	return
}

func (pm *PickleMachine) storeMemo(index int64, v interface{}) error {
	if index < 0 {
		return fmt.Errorf("Requested to write to invalid memo index:%v", index)
	}

	//If there is space in the memo presently, then store it
	//and it is done.
	if index < int64(len(pm.Memo)) {
		pm.Memo[index] = v
		return nil
	}

	//Update the maximum index in the buffer if need be
	if index > pm.memoBufferMaxDestination {
		pm.memoBufferMaxDestination = index
	}

	//If the buffer is not full write into it
	if pm.memoBufferIndex != len(pm.memoBuffer) {
		pm.memoBuffer[pm.memoBufferIndex].V = v
		pm.memoBuffer[pm.memoBufferIndex].Destination = index
		pm.memoBufferIndex++
	} else {
		//If the buffer is full flush it now
		pm.flushMemoBuffer(index, v)
	}

	return nil
}

func (pm *PickleMachine) readFromMemo(index int64) (interface{}, error) {
	if index < 0 {
		return nil, fmt.Errorf("Requested to read from negative memo index %d", index)

	}

	//Test to see if the value is outside the current length of the memo
	if index >= int64(len(pm.Memo)) {
		pm.flushMemoBuffer(-1, nil)
		if index >= int64(len(pm.Memo)) {
			return nil, fmt.Errorf("Requested to read from invalid memo index %d", index)
		}
	}

	//Grab the value
	retval := pm.Memo[index]

	//If nil then flush the memo buffer to see if it is within it
	if retval == nil {
		pm.flushMemoBuffer(-1, nil)
		//Grab the value again after the flush
		retval = pm.Memo[index]
		//If still nil, then this is a read from an invalid position
		if retval == nil {
			return nil, fmt.Errorf("Requested to read from invalid memo index %d", index)
		}
	}

	return retval, nil
}

func (pm *PickleMachine) push(v interface{}) {
	pm.Stack = append(pm.Stack, v)
}

func (pm *PickleMachine) pop() (interface{}, error) {
	l := len(pm.Stack)
	if l == 0 {
		return nil, ErrStackTooSmall
	}

	l--
	top := pm.Stack[l]
	pm.Stack = pm.Stack[:l]
	return top, nil
}

func (pm *PickleMachine) readFromStack(offset int) (interface{}, error) {
	return pm.readFromStackAt(len(pm.Stack) - 1 - offset)
}

func (pm *PickleMachine) readFromStackAt(position int) (interface{}, error) {

	if position < 0 {
		return nil, fmt.Errorf("Request to read from invalid stack position %d", position)
	}

	return pm.Stack[position], nil

}

func (pm *PickleMachine) readIntFromStack(offset int) (int64, error) {
	v, err := pm.readFromStack(offset)
	if err != nil {
		return 0, err
	}

	vi, ok := v.(int64)
	if !ok {
		return 0, fmt.Errorf("Type %T was requested from stack but found %v(%T)", vi, v, v)
	}

	return vi, nil
}

func (pm *PickleMachine) popAfterIndex(index int) {
	//Input to this function must be sane, no checking is done

	/**
	if len(pm.Stack)-1 < index {
		return ErrStackTooSmall
	}**/

	pm.Stack = pm.Stack[0:index]
}

func (pm *PickleMachine) findMark() (int, error) {
	if pm.lastMark != -1 {
		mark := pm.lastMark
		pm.lastMark = -1
		if mark < len(pm.Stack) {
			if _, ok := pm.Stack[mark].(PickleMark); ok {
				return mark, nil
			}
		}
	}

	for i := len(pm.Stack) - 1; i != -1; i-- {
		if _, ok := pm.Stack[i].(PickleMark); ok {
			return i, nil
		}
	}
	return -1, ErrMarkNotFound
}

func (pm *PickleMachine) readFixedLengthRaw(l int64) ([]byte, error) {

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return nil, err
	}

	return pm.buf.Bytes(), nil
}

func (pm *PickleMachine) readFixedLengthString(l int64) (string, error) {

	//Avoid getting "<nil>"
	if l == 0 {
		return "", nil
	}

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return "", err
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readBytes() ([]byte, error) {
	//This is slow and protocol 0 only
	pm.buf.Reset()
	for {
		var v [1]byte
		n, err := pm.Reader.Read(v[:])
		if n != 1 {
			return nil, ErrInputTruncated
		}
		if err != nil {
			return nil, err
		}

		if v[0] == '\n' {
			break
		}
		pm.buf.WriteByte(v[0])
	}

	return pm.buf.Bytes(), nil
}

func (pm *PickleMachine) readString() (string, error) {
	//This is slow and protocol 0 only
	pm.buf.Reset()
	for {
		var v [1]byte
		n, err := pm.Reader.Read(v[:])
		if n != 1 {
			return "", ErrInputTruncated
		}
		if err != nil {
			return "", err
		}

		if v[0] == '\n' {
			break
		}
		pm.buf.WriteByte(v[0])
	}

	//Avoid getting "<nil>"
	if pm.buf.Len() == 0 {
		return "", nil
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readBinaryInto(dst interface{}, bigEndian bool) error {
	var bo binary.ByteOrder
	if bigEndian {
		bo = binary.BigEndian
	} else {
		bo = binary.LittleEndian
	}
	return binary.Read(pm.Reader, bo, dst)
}