Skip to content

Commit e84325b

Browse files
committed
py: Add mp_get_stream_raise to factor out check for stream methods.
1 parent f54a96d commit e84325b

2 files changed

Lines changed: 36 additions & 40 deletions

File tree

py/stream.c

Lines changed: 29 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,21 @@ STATIC mp_obj_t stream_readall(mp_obj_t self_in);
4949

5050
#define STREAM_CONTENT_TYPE(stream) (((stream)->is_text) ? &mp_type_str : &mp_type_bytes)
5151

52-
STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
53-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(args[0]);
54-
if (o->type->stream_p == NULL || o->type->stream_p->read == NULL) {
52+
const mp_stream_p_t *mp_get_stream_raise(mp_obj_t self_in, int flags) {
53+
mp_obj_base_t *o = (mp_obj_base_t*)MP_OBJ_TO_PTR(self_in);
54+
const mp_stream_p_t *stream_p = o->type->stream_p;
55+
if (stream_p == NULL
56+
|| ((flags & MP_STREAM_OP_READ) && stream_p->read == NULL)
57+
|| ((flags & MP_STREAM_OP_WRITE) && stream_p->write == NULL)
58+
|| ((flags & MP_STREAM_OP_IOCTL) && stream_p->ioctl == NULL)) {
5559
// CPython: io.UnsupportedOperation, OSError subclass
56-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
60+
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "stream operation not supported"));
5761
}
62+
return stream_p;
63+
}
64+
65+
STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
66+
const mp_stream_p_t *stream_p = mp_get_stream_raise(args[0], MP_STREAM_OP_READ);
5867

5968
// What to do if sz < -1? Python docs don't specify this case.
6069
// CPython does a readall, but here we silently let negatives through,
@@ -65,7 +74,7 @@ STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
6574
}
6675

6776
#if MICROPY_PY_BUILTINS_STR_UNICODE
68-
if (o->type->stream_p->is_text) {
77+
if (stream_p->is_text) {
6978
// We need to read sz number of unicode characters. Because we don't have any
7079
// buffering, and because the stream API can only read bytes, we must read here
7180
// in units of bytes and must never over read. If we want sz chars, then reading
@@ -85,7 +94,7 @@ STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
8594
nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_MemoryError, "out of memory"));
8695
}
8796
int error;
88-
mp_uint_t out_sz = o->type->stream_p->read(MP_OBJ_FROM_PTR(o), p, more_bytes, &error);
97+
mp_uint_t out_sz = stream_p->read(args[0], p, more_bytes, &error);
8998
if (out_sz == MP_STREAM_ERROR) {
9099
vstr_cut_tail_bytes(&vstr, more_bytes);
91100
if (mp_is_nonblocking_error(error)) {
@@ -156,7 +165,7 @@ STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
156165
vstr_t vstr;
157166
vstr_init_len(&vstr, sz);
158167
int error;
159-
mp_uint_t out_sz = o->type->stream_p->read(MP_OBJ_FROM_PTR(o), vstr.buf, sz, &error);
168+
mp_uint_t out_sz = stream_p->read(args[0], vstr.buf, sz, &error);
160169
if (out_sz == MP_STREAM_ERROR) {
161170
vstr_clear(&vstr);
162171
if (mp_is_nonblocking_error(error)) {
@@ -170,19 +179,15 @@ STATIC mp_obj_t stream_read(mp_uint_t n_args, const mp_obj_t *args) {
170179
nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(error)));
171180
} else {
172181
vstr.len = out_sz;
173-
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(o->type->stream_p), &vstr);
182+
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(stream_p), &vstr);
174183
}
175184
}
176185

177186
mp_obj_t mp_stream_write(mp_obj_t self_in, const void *buf, size_t len) {
178-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(self_in);
179-
if (o->type->stream_p == NULL || o->type->stream_p->write == NULL) {
180-
// CPython: io.UnsupportedOperation, OSError subclass
181-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
182-
}
187+
const mp_stream_p_t *stream_p = mp_get_stream_raise(self_in, MP_STREAM_OP_WRITE);
183188

184189
int error;
185-
mp_uint_t out_sz = o->type->stream_p->write(self_in, buf, len, &error);
190+
mp_uint_t out_sz = stream_p->write(self_in, buf, len, &error);
186191
if (out_sz == MP_STREAM_ERROR) {
187192
if (mp_is_nonblocking_error(error)) {
188193
// http://docs.python.org/3/library/io.html#io.RawIOBase.write
@@ -210,11 +215,7 @@ STATIC mp_obj_t stream_write_method(mp_obj_t self_in, mp_obj_t arg) {
210215
}
211216

212217
STATIC mp_obj_t stream_readinto(mp_uint_t n_args, const mp_obj_t *args) {
213-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(args[0]);
214-
if (o->type->stream_p == NULL || o->type->stream_p->read == NULL) {
215-
// CPython: io.UnsupportedOperation, OSError subclass
216-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
217-
}
218+
const mp_stream_p_t *stream_p = mp_get_stream_raise(args[0], MP_STREAM_OP_READ);
218219
mp_buffer_info_t bufinfo;
219220
mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_WRITE);
220221

@@ -230,7 +231,7 @@ STATIC mp_obj_t stream_readinto(mp_uint_t n_args, const mp_obj_t *args) {
230231
}
231232

232233
int error;
233-
mp_uint_t out_sz = o->type->stream_p->read(MP_OBJ_FROM_PTR(o), bufinfo.buf, len, &error);
234+
mp_uint_t out_sz = stream_p->read(args[0], bufinfo.buf, len, &error);
234235
if (out_sz == MP_STREAM_ERROR) {
235236
if (mp_is_nonblocking_error(error)) {
236237
return mp_const_none;
@@ -242,11 +243,7 @@ STATIC mp_obj_t stream_readinto(mp_uint_t n_args, const mp_obj_t *args) {
242243
}
243244

244245
STATIC mp_obj_t stream_readall(mp_obj_t self_in) {
245-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(self_in);
246-
if (o->type->stream_p == NULL || o->type->stream_p->read == NULL) {
247-
// CPython: io.UnsupportedOperation, OSError subclass
248-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
249-
}
246+
const mp_stream_p_t *stream_p = mp_get_stream_raise(self_in, MP_STREAM_OP_READ);
250247

251248
mp_uint_t total_size = 0;
252249
vstr_t vstr;
@@ -255,7 +252,7 @@ STATIC mp_obj_t stream_readall(mp_obj_t self_in) {
255252
mp_uint_t current_read = DEFAULT_BUFFER_SIZE;
256253
while (true) {
257254
int error;
258-
mp_uint_t out_sz = o->type->stream_p->read(self_in, p, current_read, &error);
255+
mp_uint_t out_sz = stream_p->read(self_in, p, current_read, &error);
259256
if (out_sz == MP_STREAM_ERROR) {
260257
if (mp_is_nonblocking_error(error)) {
261258
// With non-blocking streams, we read as much as we can.
@@ -286,16 +283,12 @@ STATIC mp_obj_t stream_readall(mp_obj_t self_in) {
286283
}
287284

288285
vstr.len = total_size;
289-
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(o->type->stream_p), &vstr);
286+
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(stream_p), &vstr);
290287
}
291288

292289
// Unbuffered, inefficient implementation of readline() for raw I/O files.
293290
STATIC mp_obj_t stream_unbuffered_readline(mp_uint_t n_args, const mp_obj_t *args) {
294-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(args[0]);
295-
if (o->type->stream_p == NULL || o->type->stream_p->read == NULL) {
296-
// CPython: io.UnsupportedOperation, OSError subclass
297-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
298-
}
291+
const mp_stream_p_t *stream_p = mp_get_stream_raise(args[0], MP_STREAM_OP_READ);
299292

300293
mp_int_t max_size = -1;
301294
if (n_args > 1) {
@@ -316,7 +309,7 @@ STATIC mp_obj_t stream_unbuffered_readline(mp_uint_t n_args, const mp_obj_t *arg
316309
}
317310

318311
int error;
319-
mp_uint_t out_sz = o->type->stream_p->read(MP_OBJ_FROM_PTR(o), p, 1, &error);
312+
mp_uint_t out_sz = stream_p->read(args[0], p, 1, &error);
320313
if (out_sz == MP_STREAM_ERROR) {
321314
if (mp_is_nonblocking_error(error)) {
322315
if (vstr.len == 1) {
@@ -347,7 +340,7 @@ STATIC mp_obj_t stream_unbuffered_readline(mp_uint_t n_args, const mp_obj_t *arg
347340
}
348341
}
349342

350-
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(o->type->stream_p), &vstr);
343+
return mp_obj_new_str_from_vstr(STREAM_CONTENT_TYPE(stream_p), &vstr);
351344
}
352345

353346
// TODO take an optional extra argument (what does it do exactly?)
@@ -373,11 +366,7 @@ mp_obj_t mp_stream_unbuffered_iter(mp_obj_t self) {
373366
}
374367

375368
STATIC mp_obj_t stream_seek(mp_uint_t n_args, const mp_obj_t *args) {
376-
struct _mp_obj_base_t *o = (struct _mp_obj_base_t *)MP_OBJ_TO_PTR(args[0]);
377-
if (o->type->stream_p == NULL || o->type->stream_p->ioctl == NULL) {
378-
// CPython: io.UnsupportedOperation, OSError subclass
379-
nlr_raise(mp_obj_new_exception_msg(&mp_type_OSError, "Operation not supported"));
380-
}
369+
const mp_stream_p_t *stream_p = mp_get_stream_raise(args[0], MP_STREAM_OP_IOCTL);
381370

382371
struct mp_stream_seek_t seek_s;
383372
// TODO: Could be uint64
@@ -388,7 +377,7 @@ STATIC mp_obj_t stream_seek(mp_uint_t n_args, const mp_obj_t *args) {
388377
}
389378

390379
int error;
391-
mp_uint_t res = o->type->stream_p->ioctl(MP_OBJ_FROM_PTR(o), MP_STREAM_SEEK, (mp_uint_t)(uintptr_t)&seek_s, &error);
380+
mp_uint_t res = stream_p->ioctl(args[0], MP_STREAM_SEEK, (mp_uint_t)(uintptr_t)&seek_s, &error);
392381
if (res == MP_STREAM_ERROR) {
393382
nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(error)));
394383
}

py/stream.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,13 @@ MP_DECLARE_CONST_FUN_OBJ(mp_stream_write_obj);
3737
MP_DECLARE_CONST_FUN_OBJ(mp_stream_seek_obj);
3838
MP_DECLARE_CONST_FUN_OBJ(mp_stream_tell_obj);
3939

40+
// these are for mp_get_stream_raise and can be or'd together
41+
#define MP_STREAM_OP_READ (1)
42+
#define MP_STREAM_OP_WRITE (2)
43+
#define MP_STREAM_OP_IOCTL (4)
44+
45+
const mp_stream_p_t *mp_get_stream_raise(mp_obj_t self_in, int flags);
46+
4047
// Iterator which uses mp_stream_unbuffered_readline_obj
4148
mp_obj_t mp_stream_unbuffered_iter(mp_obj_t self);
4249

0 commit comments

Comments
 (0)