DOCSEARCHbyAlgolia
Submit Docs

GoPickle

repository·master·Indexed 41 days ago

https://github.com/nlpodyssey/gopickle

A Go library for loading Python-serialized data, supporting Python pickle protocols 0-5 and PyTorch module files (.pt). It features a virtual unpickling machine and a pytorch sub-package for handling modern zip-compressed and legacy non-tar formats.

Tokens
1.2K
Snippets
3
Records
4
Agent score
50%

What's inside GoPickle

  1. How GoPickle's unpickling machine works

    master

    GoPickle implements a virtual unpickling machine (PM) that executes a sequence of opcodes to reconstruct Python objects.

    Core Components

    • Opcodes: The machine processes a sequence of instructions (similar to Python's pickletools) to build complex objects.
    • Types: The library uses a types sub-package containing Go representations of Python classes, instances, and common builtin non-scalar types.
    • Extensibility: Because Python can serialize any object, GoPickle uses a GenericClass for unknown types. If a pickle requires specific traits that GenericClass cannot provide, users must implement a custom class via the FindClass callback.
    • Interface-based design: Once objects are resolved, they are cast to interface{}. The machine only checks for specific interfaces when strictly necessary to maintain performance and simplicity.

  2. Load Python pickle data with the pickle sub-package

    master

    The pickle sub-package allows you to load data serialized with Python's pickle module from files or strings. It supports all pickle protocols from 0 to 5.

    Simple Loading

    Use pickle.Load to read from a file or pickle.Loads to read from a string.

    Advanced Customization

    For complex pickles that require custom class mapping, persistent IDs, or specific protocol extensions, use pickle.NewUnpickler(io.Reader). You can provide several callback functions to the Unpickler to handle specialized Python objects.

    import "github.com/nlpodyssey/gopickle/pickle"

// from file
foo, err := pickle.Load("foo.p") 

// from string
stringDump := "I42\n."
bar, err := pickle.Loads(stringDump)

  3. Load PyTorch module files

    master

    The pytorch sub-package provides high-level functions to load PyTorch module files (.pt).

    Supported formats:

    • Modern zip-compressed format
    • Legacy non-tar format

    Unsupported formats:

    • Legacy tar-compressed files
    • TorchScript archives
    import "github.com/nlpodyssey/gopickle/pytorch"

// Load a PyTorch module file
myModel, err := pytorch.Load("module.pt")

  4. Customize Unpickler behavior with callbacks

    master

    When using pickle.NewUnpickler(r), you can override default behaviors by assigning functions to the following fields on the Unpickler object:

    • FindClass: Maps a module and name to a Go implementation of a Python class. If not provided, a GenericClass is used by default.
    • PersistentLoad: Resolves objects by a persistentId.
    • GetExtension: Handles custom pickle extensions via an integer code.
    • NextBuffer: Handles Out-of-band Buffers (Python 3 pickle protocol).
    • MakeReadOnly: A low-level handler for the protocol 5 READONLY_BUFFER opcode, allowing you to transform objects into a read-only state.
    import (
    "fmt"
    "io"
    "github.com/nlpodyssey/gopickle/pickle"
)

var r io.Reader
u := pickle.NewUnpickler(r)

// Handle custom classes
u.FindClass = func(module, name string) (interface{}, error) {
    if module == "foo" && name == "Bar" {
        return myFooBarClass, nil
    }
    return nil, fmt.Errorf("class not found :(")
}

// Resolve objects by persistent ID
u.PersistentLoad = func(persistentId interface{}) (interface{}, error) {
    obj := doSomethingWithPersistentId(persistentId)
    return obj, nil
}

// Handle custom pickle extensions
u.GetExtension = func(code int) (interface{}, error) {
    obj := doSomethingToResolveExtension(code)
    return obj, nil
}

// Handle Out-of-band Buffers
u.NextBuffer = func() (interface{}, error) {
    buf := getMyNextBuffer()
    return buf, nil
}

// Handle protocol 5 READONLY_BUFFER opcode
u.MakeReadOnly = func(obj interface{}) (interface{}, error) {
    newObj := myReadOnlyTransform(obj)
    return newObj, nil
}

data, err := u.Load()