40
40
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
41
41
"""
42
42
43
+ import os
43
44
from datetime import datetime , date , timedelta
44
45
from dateutil .parser import parse
45
46
54
55
from pandas .sparse .array import BlockIndex , IntIndex
55
56
from pandas .core .generic import NDFrame
56
57
from pandas .core .common import needs_i8_conversion
58
+ from pandas .io .common import get_filepath_or_buffer
57
59
from pandas .core .internals import BlockManager , make_block
58
60
import pandas .core .internals as internals
59
61
71
73
compressor = None
72
74
73
75
74
- def to_msgpack (path , * args , ** kwargs ):
76
+ def to_msgpack (path_or_buf , * args , ** kwargs ):
75
77
"""
76
78
msgpack (serialize) object to input file path
77
79
@@ -80,7 +82,8 @@ def to_msgpack(path, *args, **kwargs):
80
82
81
83
Parameters
82
84
----------
83
- path : string File path
85
+ path_or_buf : string File path, buffer-like, or None
86
+ if None, return generated string
84
87
args : an object or objects to serialize
85
88
append : boolean whether to append to an existing msgpack
86
89
(default is False)
@@ -90,17 +93,24 @@ def to_msgpack(path, *args, **kwargs):
90
93
compressor = kwargs .pop ('compress' , None )
91
94
append = kwargs .pop ('append' , None )
92
95
if append :
93
- f = open ( path , 'a+b' )
96
+ mode = 'a+b'
94
97
else :
95
- f = open (path , 'wb' )
96
- try :
97
- for a in args :
98
- f .write (pack (a , ** kwargs ))
99
- finally :
100
- f .close ()
98
+ mode = 'wb'
101
99
100
+ def writer (fh ):
101
+ for a in args :
102
+ fh .write (pack (a , ** kwargs ))
103
+ return fh
104
+
105
+ if isinstance (path_or_buf , compat .string_types ):
106
+ with open (path_or_buf , mode ) as fh :
107
+ writer (fh )
108
+ elif path_or_buf is None :
109
+ return writer (compat .BytesIO ())
110
+ else :
111
+ writer (path_or_buf )
102
112
103
- def read_msgpack (path , iterator = False , ** kwargs ):
113
+ def read_msgpack (path_or_buf , iterator = False , ** kwargs ):
104
114
"""
105
115
Load msgpack pandas object from the specified
106
116
file path
@@ -110,8 +120,7 @@ def read_msgpack(path, iterator=False, **kwargs):
110
120
111
121
Parameters
112
122
----------
113
- path : string
114
- File path
123
+ path_or_buf : string File path, BytesIO like or string
115
124
iterator : boolean, if True, return an iterator to the unpacker
116
125
(default is False)
117
126
@@ -120,15 +129,40 @@ def read_msgpack(path, iterator=False, **kwargs):
120
129
obj : type of object stored in file
121
130
122
131
"""
132
+ path_or_buf , _ = get_filepath_or_buffer (path_or_buf )
123
133
if iterator :
124
- return Iterator (path )
134
+ return Iterator (path_or_buf )
125
135
126
- with open ( path , 'rb' ) as fh :
136
+ def read ( fh ) :
127
137
l = list (unpack (fh ))
128
138
if len (l ) == 1 :
129
139
return l [0 ]
130
140
return l
131
141
142
+ # see if we have an actual file
143
+ if isinstance (path_or_buf , compat .string_types ):
144
+
145
+ try :
146
+ path_exists = os .path .exists (path_or_buf )
147
+ except (TypeError ):
148
+ path_exists = False
149
+
150
+ if path_exists :
151
+ with open (path_or_buf , 'rb' ) as fh :
152
+ return read (fh )
153
+
154
+ # treat as a string-like
155
+ if not hasattr (path_or_buf ,'read' ):
156
+
157
+ try :
158
+ fh = compat .BytesIO (path_or_buf )
159
+ return read (fh )
160
+ finally :
161
+ fh .close ()
162
+
163
+ # a buffer like
164
+ return read (path_or_buf )
165
+
132
166
dtype_dict = {21 : np .dtype ('M8[ns]' ),
133
167
u ('datetime64[ns]' ): np .dtype ('M8[ns]' ),
134
168
u ('datetime64[us]' ): np .dtype ('M8[us]' ),
@@ -530,10 +564,36 @@ def __init__(self, path, **kwargs):
530
564
531
565
def __iter__ (self ):
532
566
567
+ needs_closing = True
533
568
try :
534
- fh = open (self .path , 'rb' )
569
+
570
+ # see if we have an actual file
571
+ if isinstance (self .path , compat .string_types ):
572
+
573
+ try :
574
+ path_exists = os .path .exists (self .path )
575
+ except (TypeError ):
576
+ path_exists = False
577
+
578
+ if path_exists :
579
+ fh = open (self .path , 'rb' )
580
+ else :
581
+ fh = compat .BytesIO (self .path )
582
+
583
+ else :
584
+
585
+ if not hasattr (self .path ,'read' ):
586
+ fh = compat .BytesIO (self .path )
587
+
588
+ else :
589
+
590
+ # a file-like
591
+ needs_closing = False
592
+ fh = self .path
593
+
535
594
unpacker = unpack (fh )
536
595
for o in unpacker :
537
596
yield o
538
597
finally :
539
- fh .close ()
598
+ if needs_closing :
599
+ fh .close ()
0 commit comments