1/*
2 * This file is part of the MicroPython project, http://micropython.org/
3 *
4 * The MIT License (MIT)
5 *
6 * Copyright (c) 2013, 2014 Damien P. George
7 * Copyright (c) 2014-2017 Paul Sokolovsky
8 *
9 * Permission is hereby granted, free of charge, to any person obtaining a copy
10 * of this software and associated documentation files (the "Software"), to deal
11 * in the Software without restriction, including without limitation the rights
12 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13 * copies of the Software, and to permit persons to whom the Software is
14 * furnished to do so, subject to the following conditions:
15 *
16 * The above copyright notice and this permission notice shall be included in
17 * all copies or substantial portions of the Software.
18 *
19 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
25 * THE SOFTWARE.
26 */
27
28#include <stdio.h>
29#include <string.h>
30
31#include "py/objstr.h"
32#include "py/objstringio.h"
33#include "py/runtime.h"
34#include "py/stream.h"
35
36#if MICROPY_PY_IO
37
38#if MICROPY_CPYTHON_COMPAT
39STATIC void check_stringio_is_open(const mp_obj_stringio_t *o) {
40 if (o->vstr == NULL) {
41 mp_raise_ValueError(MP_ERROR_TEXT("I/O operation on closed file"));
42 }
43}
44#else
45#define check_stringio_is_open(o)
46#endif
47
48STATIC void stringio_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
49 (void)kind;
50 mp_obj_stringio_t *self = MP_OBJ_TO_PTR(self_in);
51 mp_printf(print, self->base.type == &mp_type_stringio ? "<io.StringIO 0x%x>" : "<io.BytesIO 0x%x>", self);
52}
53
54STATIC mp_uint_t stringio_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *errcode) {
55 (void)errcode;
56 mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in);
57 check_stringio_is_open(o);
58 if (o->vstr->len <= o->pos) { // read to EOF, or seeked to EOF or beyond
59 return 0;
60 }
61 mp_uint_t remaining = o->vstr->len - o->pos;
62 if (size > remaining) {
63 size = remaining;
64 }
65 memcpy(buf, o->vstr->buf + o->pos, size);
66 o->pos += size;
67 return size;
68}
69
70STATIC void stringio_copy_on_write(mp_obj_stringio_t *o) {
71 const void *buf = o->vstr->buf;
72 o->vstr->buf = m_new(char, o->vstr->len);
73 o->vstr->fixed_buf = false;
74 o->ref_obj = MP_OBJ_NULL;
75 memcpy(o->vstr->buf, buf, o->vstr->len);
76}
77
78STATIC mp_uint_t stringio_write(mp_obj_t o_in, const void *buf, mp_uint_t size, int *errcode) {
79 (void)errcode;
80 mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in);
81 check_stringio_is_open(o);
82
83 if (o->vstr->fixed_buf) {
84 stringio_copy_on_write(o);
85 }
86
87 mp_uint_t new_pos = o->pos + size;
88 if (new_pos < size) {
89 // Writing <size> bytes will overflow o->pos beyond limit of mp_uint_t.
90 *errcode = MP_EFBIG;
91 return MP_STREAM_ERROR;
92 }
93 mp_uint_t org_len = o->vstr->len;
94 if (new_pos > o->vstr->alloc) {
95 // Take all what's already allocated...
96 o->vstr->len = o->vstr->alloc;
97 // ... and add more
98 vstr_add_len(o->vstr, new_pos - o->vstr->alloc);
99 }
100 // If there was a seek past EOF, clear the hole
101 if (o->pos > org_len) {
102 memset(o->vstr->buf + org_len, 0, o->pos - org_len);
103 }
104 memcpy(o->vstr->buf + o->pos, buf, size);
105 o->pos = new_pos;
106 if (new_pos > o->vstr->len) {
107 o->vstr->len = new_pos;
108 }
109 return size;
110}
111
112STATIC mp_uint_t stringio_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg, int *errcode) {
113 (void)errcode;
114 mp_obj_stringio_t *o = MP_OBJ_TO_PTR(o_in);
115 switch (request) {
116 case MP_STREAM_SEEK: {
117 struct mp_stream_seek_t *s = (struct mp_stream_seek_t *)arg;
118 mp_uint_t ref = 0;
119 switch (s->whence) {
120 case MP_SEEK_CUR:
121 ref = o->pos;
122 break;
123 case MP_SEEK_END:
124 ref = o->vstr->len;
125 break;
126 }
127 mp_uint_t new_pos = ref + s->offset;
128
129 // For MP_SEEK_SET, offset is unsigned
130 if (s->whence != MP_SEEK_SET && s->offset < 0) {
131 if (new_pos > ref) {
132 // Negative offset from SEEK_CUR or SEEK_END went past 0.
133 // CPython sets position to 0, POSIX returns an EINVAL error
134 new_pos = 0;
135 }
136 } else if (new_pos < ref) {
137 // positive offset went beyond the limit of mp_uint_t
138 *errcode = MP_EINVAL; // replace with MP_EOVERFLOW when defined
139 return MP_STREAM_ERROR;
140 }
141 s->offset = o->pos = new_pos;
142 return 0;
143 }
144 case MP_STREAM_FLUSH:
145 return 0;
146 case MP_STREAM_CLOSE:
147 #if MICROPY_CPYTHON_COMPAT
148 vstr_free(o->vstr);
149 o->vstr = NULL;
150 #else
151 vstr_clear(o->vstr);
152 o->vstr->alloc = 0;
153 o->vstr->len = 0;
154 o->pos = 0;
155 #endif
156 return 0;
157 default:
158 *errcode = MP_EINVAL;
159 return MP_STREAM_ERROR;
160 }
161}
162
163#define STREAM_TO_CONTENT_TYPE(o) (((o)->base.type == &mp_type_stringio) ? &mp_type_str : &mp_type_bytes)
164
165STATIC mp_obj_t stringio_getvalue(mp_obj_t self_in) {
166 mp_obj_stringio_t *self = MP_OBJ_TO_PTR(self_in);
167 check_stringio_is_open(self);
168 // TODO: Try to avoid copying string
169 return mp_obj_new_str_of_type(STREAM_TO_CONTENT_TYPE(self), (byte *)self->vstr->buf, self->vstr->len);
170}
171STATIC MP_DEFINE_CONST_FUN_OBJ_1(stringio_getvalue_obj, stringio_getvalue);
172
173STATIC mp_obj_t stringio___exit__(size_t n_args, const mp_obj_t *args) {
174 (void)n_args;
175 return mp_stream_close(args[0]);
176}
177STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(stringio___exit___obj, 4, 4, stringio___exit__);
178
179STATIC mp_obj_stringio_t *stringio_new(const mp_obj_type_t *type) {
180 mp_obj_stringio_t *o = m_new_obj(mp_obj_stringio_t);
181 o->base.type = type;
182 o->pos = 0;
183 o->ref_obj = MP_OBJ_NULL;
184 return o;
185}
186
187STATIC mp_obj_t stringio_make_new(const mp_obj_type_t *type_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
188 (void)n_kw; // TODO check n_kw==0
189
190 mp_uint_t sz = 16;
191 bool initdata = false;
192 mp_buffer_info_t bufinfo;
193
194 mp_obj_stringio_t *o = stringio_new(type_in);
195
196 if (n_args > 0) {
197 if (mp_obj_is_int(args[0])) {
198 sz = mp_obj_get_int(args[0]);
199 } else {
200 mp_get_buffer_raise(args[0], &bufinfo, MP_BUFFER_READ);
201
202 if (mp_obj_is_str_or_bytes(args[0])) {
203 o->vstr = m_new_obj(vstr_t);
204 vstr_init_fixed_buf(o->vstr, bufinfo.len, bufinfo.buf);
205 o->vstr->len = bufinfo.len;
206 o->ref_obj = args[0];
207 return MP_OBJ_FROM_PTR(o);
208 }
209
210 sz = bufinfo.len;
211 initdata = true;
212 }
213 }
214
215 o->vstr = vstr_new(sz);
216
217 if (initdata) {
218 stringio_write(MP_OBJ_FROM_PTR(o), bufinfo.buf, bufinfo.len, NULL);
219 // Cur ptr is always at the beginning of buffer at the construction
220 o->pos = 0;
221 }
222 return MP_OBJ_FROM_PTR(o);
223}
224
225STATIC const mp_rom_map_elem_t stringio_locals_dict_table[] = {
226 { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
227 { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
228 { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
229 { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
230 { MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
231 { MP_ROM_QSTR(MP_QSTR_tell), MP_ROM_PTR(&mp_stream_tell_obj) },
232 { MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&mp_stream_flush_obj) },
233 { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
234 { MP_ROM_QSTR(MP_QSTR_getvalue), MP_ROM_PTR(&stringio_getvalue_obj) },
235 { MP_ROM_QSTR(MP_QSTR___enter__), MP_ROM_PTR(&mp_identity_obj) },
236 { MP_ROM_QSTR(MP_QSTR___exit__), MP_ROM_PTR(&stringio___exit___obj) },
237};
238
239STATIC MP_DEFINE_CONST_DICT(stringio_locals_dict, stringio_locals_dict_table);
240
241STATIC const mp_stream_p_t stringio_stream_p = {
242 .read = stringio_read,
243 .write = stringio_write,
244 .ioctl = stringio_ioctl,
245 .is_text = true,
246};
247
248const mp_obj_type_t mp_type_stringio = {
249 { &mp_type_type },
250 .name = MP_QSTR_StringIO,
251 .print = stringio_print,
252 .make_new = stringio_make_new,
253 .getiter = mp_identity_getiter,
254 .iternext = mp_stream_unbuffered_iter,
255 .protocol = &stringio_stream_p,
256 .locals_dict = (mp_obj_dict_t *)&stringio_locals_dict,
257};
258
259#if MICROPY_PY_IO_BYTESIO
260STATIC const mp_stream_p_t bytesio_stream_p = {
261 .read = stringio_read,
262 .write = stringio_write,
263 .ioctl = stringio_ioctl,
264};
265
266const mp_obj_type_t mp_type_bytesio = {
267 { &mp_type_type },
268 .name = MP_QSTR_BytesIO,
269 .print = stringio_print,
270 .make_new = stringio_make_new,
271 .getiter = mp_identity_getiter,
272 .iternext = mp_stream_unbuffered_iter,
273 .protocol = &bytesio_stream_p,
274 .locals_dict = (mp_obj_dict_t *)&stringio_locals_dict,
275};
276#endif
277
278#endif
279