gt_io.py 2.33 KB
Newer Older
1
#! /usr/bin/env python
2
# -*- coding: utf-8 -*-
3
#
4
5
# graph_tool -- a general graph manipulation python module
#
Tiago Peixoto's avatar
Tiago Peixoto committed
6
# Copyright (C) 2006-2020 Tiago de Paula Peixoto <tiago@skewed.de>
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

21
import pickle
22
import atexit
23
from io import BytesIO
24
from . import libgraph_tool_core
25
26
27
28

# IStream and OStream need to be tweaked a little to become a real file-like
# object...

29
30
31
32
def IStream_read(self, n=None, buflen=1048576):
    if n is None:
        data = b""
        while True:
33
            buf = self.read_buf(buflen)
34
35
36
            data += buf
            if len(buf) < buflen:
                break
37
38
        return data
    else:
39
        return self.read_buf(n)
Tiago Peixoto's avatar
Tiago Peixoto committed
40
41

def IStream_readline(self, n=None):
42
    c = None
43
    line = b""
44
    while c != "" and c != "\n" and len(line) < n:
45
        c = self.read_buf(1)
46
47
48
49
        line += c
    return line

def OStream_write(self, s):
50
    data = s
51
    self.write_buf(data, len(s))
52

53
54
libgraph_tool_core.IStream.read_buf = libgraph_tool_core.IStream.read
libgraph_tool_core.OStream.write_buf = libgraph_tool_core.OStream.write
55
56
57
58
59
libgraph_tool_core.IStream.read = IStream_read
libgraph_tool_core.IStream.readline = IStream_readline
libgraph_tool_core.OStream.write = OStream_write

# define and set the pickler/unpickler functions
60
61
62

GT_PICKLE_PROTOCOL = -1

63
def pickler(stream, obj):
64
    sstream = BytesIO()
65
    pickle.dump(obj, sstream, GT_PICKLE_PROTOCOL)
66
    stream.write(sstream.getvalue())
Tiago Peixoto's avatar
Tiago Peixoto committed
67

68
def unpickler(stream):
69
70
    data = stream.read()
    sstream = BytesIO(data)
71
    return pickle.load(sstream, encoding="bytes")
72
73
74

libgraph_tool_core.set_pickler(pickler)
libgraph_tool_core.set_unpickler(unpickler)
75
76
77
78
79
80

def clean_picklers():
    libgraph_tool_core.set_pickler(None)
    libgraph_tool_core.set_unpickler(None)

atexit.register(clean_picklers)