gt_io.py 2.45 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-2016 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
22
from __future__ import division, absolute_import, print_function

23
import pickle
Tiago Peixoto's avatar
Tiago Peixoto committed
24
import base64
25
import atexit
26
import sys
27
from io import BytesIO
28
from . import libgraph_tool_core
29
30
31
32

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

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

def IStream_readline(self, n=None):
46
    c = None
47
    line = b""
48
    while c != "" and c != "\n" and len(line) < n:
49
        c = self.read_buf(1)
50
51
52
53
        line += c
    return line

def OStream_write(self, s):
54
    data = s
55
    self.write_buf(data, len(s))
56

57
58
libgraph_tool_core.IStream.read_buf = libgraph_tool_core.IStream.read
libgraph_tool_core.OStream.write_buf = libgraph_tool_core.OStream.write
59
60
61
62
63
64
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
def pickler(stream, obj):
65
    sstream = BytesIO()
66
    pickle.dump(obj, sstream, -1)
67
    stream.write(sstream.getvalue())
Tiago Peixoto's avatar
Tiago Peixoto committed
68

69
def unpickler(stream):
70
71
    data = stream.read()
    sstream = BytesIO(data)
72
73
74
    if sys.version_info < (3,):
        return pickle.load(sstream)
    return pickle.load(sstream, encoding="bytes")
75
76
77

libgraph_tool_core.set_pickler(pickler)
libgraph_tool_core.set_unpickler(unpickler)
78
79
80
81
82
83

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

atexit.register(clean_picklers)