io.py 2.2 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) 2007-2012 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
from io import BytesIO
27
from . import libgraph_tool_core
28
29
30
31

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

Tiago Peixoto's avatar
Tiago Peixoto committed
32
33

def IStream_read(self, n=None):
34
35
36
37
38
39
40
41
42
43
    if n == None:
        data = ""
        new_data = None
        while new_data != "":
            new_data = self.Read(1)
            data += new_data
        return data
    else:
        return self.Read(n)

Tiago Peixoto's avatar
Tiago Peixoto committed
44
45

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

Tiago Peixoto's avatar
Tiago Peixoto committed
53

54
55
56
57
58
59
60
def OStream_write(self, s):
    self.Write(s, len(s))

libgraph_tool_core.IStream.read = IStream_read
libgraph_tool_core.IStream.readline = IStream_readline
libgraph_tool_core.OStream.write = OStream_write

Tiago Peixoto's avatar
Tiago Peixoto committed
61

62
63
# define and set the pickler/unpickler functions
def pickler(stream, obj):
64
    sstream = BytesIO()
65
    pickle.dump(obj, sstream, -1)
Tiago Peixoto's avatar
Tiago Peixoto committed
66
    stream.write(base64.b64encode(sstream.getvalue().encode('ascii')))
67

Tiago Peixoto's avatar
Tiago Peixoto committed
68

69
def unpickler(stream):
Tiago Peixoto's avatar
Tiago Peixoto committed
70
    sstream = BytesIO(base64.b64decode(stream.read().encode('ascii')))
71
    return pickle.load(sstream)
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)