8000 py/stream.c: Implemented truncate(). · micropython/micropython@ec18976 · GitHub
[go: up one dir, main page]
Skip to content

Commit ec18976

Browse files
DavidEGraysonDavidEGrayson
committed
py/stream.c: Implemented truncate().
Signed-off-by: David Grayson <davidegrayson@gmail.com>
1 parent 3229791 commit ec18976

File tree

7 files changed

+107
-1
lines changed

7 files changed

+107
-1
lines changed

‎extmod/vfs_fat_file.c

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,20 @@ STATIC mp_uint_t file_obj_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg,
126126
s->offset = f_tell(&self->fp);
127127
return 0;
128128

129+
} else if (request == MP_STREAM_TRUNCATE) {
130+
// f_truncate does not take a size argument so we have to do extra work.
131+
uint32_t pos = f_tell(&self->fp);
132+
f_lseek(&self->fp, arg);
133+
FRESULT res = f_truncate(&self->fp);
134+
if (pos < arg) {
135+
f_lseek(&self->fp, pos);
136+
}
137+
if (res != FR_OK) {
138+
*errcode = fresult_to_errno_table[res];
139+
return MP_STREAM_ERROR;
140+
}
141+
return 0;
142+
129143
} else if (request == MP_STREAM_FLUSH) {
130144
FRESULT res = f_sync(&self->fp);
131145
if (res != FR_OK) {
@@ -144,7 +158,6 @@ STATIC mp_uint_t file_obj_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg,
144158
}
145159
}
146160
return 0;
147-
148161
} else {
149162
*errcode = MP_EINVAL;
150163
return MP_STREAM_ERROR;
@@ -163,6 +176,7 @@ STATIC const mp_rom_map_elem_t vfs_fat_rawfile_locals_dict_table[] = {
163176
{ MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
164177
{ MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
165178
{ MP_ROM_QSTR(MP_QSTR_tell), MP_ROM_PTR(&mp_stream_tell_obj) },
179+
{ MP_ROM_QSTR(MP_QSTR_truncate), MP_ROM_PTR(&mp_stream_truncate_obj) },
166180
{ MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&mp_stream_close_obj) },
167181
{ MP_ROM_QSTR(MP_QSTR___enter__), MP_ROM_PTR(&mp_identity_obj) },
168182
{ MP_ROM_QSTR(MP_QSTR___exit__), MP_ROM_PTR(&file_obj___exit___obj) },

‎extmod/vfs_lfsx_file.c

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,13 @@ STATIC mp_uint_t MP_VFS_LFSx(file_ioctl)(mp_obj_t self_in, mp_uint_t request, ui
177177
}
178178
s->offset = res;
179179
return 0;
180+
} else if (request == MP_STREAM_TRUNCATE) {
181+
int res = LFSx_API(file_truncate)(&self->vfs->lfs, &self->file, arg);
182+
if (res < 0) {
183+
*errcode = -res;
184+
return MP_STREAM_ERROR;
185+
}
186+
return 0;
180187
} else if (request == MP_STREAM_FLUSH) {
181188
int res = LFSx_API(file_sync)(&self->vfs->lfs, &self->file);
182189
if (res < 0) {
@@ -211,6 +218,7 @@ STATIC const mp_rom_map_elem_t MP_VFS_LFSx(file_locals_dict_table)[] = {
211218
{ MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
212219
{ MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
213220
{ MP_ROM_QSTR(MP_QSTR_tell), MP_ROM_PTR(&mp_stream_tell_obj) },
221+
{ MP_ROM_QSTR(MP_QSTR_truncate), MP_ROM_PTR(&mp_stream_truncate_obj) },
214222
{ MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&mp_stream_close_obj) },
215223
{ MP_ROM_QSTR(MP_QSTR___enter__), MP_ROM_PTR(&mp_identity_obj) },
216224
{ MP_ROM_QSTR(MP_QSTR___exit__), MP_ROM_PTR(&MP_VFS_LFSx(file___exit___obj)) },

‎extmod/vfs_posix_file.c

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,16 @@ STATIC mp_uint_t vfs_posix_file_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_
184184
s->offset = off;
185185
return 0;
186186
}
187+
case MP_STREAM_TRUNCATE: {
188+
MP_THREAD_GIL_EXIT();
189+
int res = ftruncate(o->fd, arg);
190+
MP_THREAD_GIL_ENTER();
191+
if (res < 0) {
192+
*errcode = errno;
193+
return MP_STREAM_ERROR;
194+
}
195+
return 0;
196+
}
187197
case MP_STREAM_CLOSE:
188198
if (o->fd >= 0) {
189199
MP_THREAD_GIL_EXIT();
@@ -235,6 +245,7 @@ STATIC const mp_rom_map_elem_t vfs_posix_rawfile_locals_dict_table[] = {
235245
{ MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
236246
{ MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
237247
{ MP_ROM_QSTR(MP_QSTR_tell), MP_ROM_PTR(&mp_stream_tell_obj) },
248+
{ MP_ROM_QSTR(MP_QSTR_truncate), MP_ROM_PTR(&mp_stream_truncate_obj) },
238249
{ MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&mp_stream_flush_obj) },
239250
{ MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
240251
{ MP_ROM_QSTR(MP_QSTR___del__), MP_ROM_PTR(&mp_stream_close_obj) },

‎py/objstringio.c

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,15 @@ STATIC mp_uint_t stringio_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg,
141141
s->offset = o->pos = new_pos;
142142
return 0;
143143
}
144+
case MP_STREAM_TRUNCATE:
145+
// In CPython, StringIO.truncate never expands the string.
146+
if (o->vstr->len > arg) {
147+
o->vstr->len = arg;
148+
if (o->pos > arg) {
149+
o->pos = arg;
150+
}
151+
}
152+
return 0;
144153
case MP_STREAM_FLUSH:
145154
return 0;
146155
case MP_STREAM_CLOSE:
@@ -228,6 +237,7 @@ STATIC const mp_rom_map_elem_t stringio_locals_dict_table[] = {
228237
{ MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
229238
{ MP_ROM_QSTR(MP_QSTR_seek), MP_ROM_PTR(&mp_stream_seek_obj) },
230239
{ MP_ROM_QSTR(MP_QSTR_tell), MP_ROM_PTR(&mp_stream_tell_obj) },
240+
{ MP_ROM_QSTR(MP_QSTR_truncate), MP_ROM_PTR(&mp_stream_truncate_obj) },
231241
{ MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&mp_stream_flush_obj) },
232242
{ MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
233243
{ MP_ROM_QSTR(MP_QSTR_getvalue), MP_ROM_PTR(&stringio_getvalue_obj) },

‎py/stream.c

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,27 @@ STATIC mp_obj_t stream_tell(mp_obj_t self) {
472472
}
473473
MP_DEFINE_CONST_FUN_OBJ_1(mp_stream_tell_obj, stream_tell);
474474

475+
STATIC mp_obj_t stream_truncate(size_t n_args, const mp_obj_t *args) {
476+
const mp_stream_p_t *stream_p = mp_get_stream(args[0]);
477+
mp_uint_t size;
478+
if (n_args > 1) {
479+
mp_int_t arg = mp_obj_get_int(args[1]);
480+
if (arg < 0) {
481+
mp_raise_OSError(MP_EINVAL);
482+
}
483+
size = arg;
484+
} else {
485+
size = mp_obj_get_int(stream_tell(args[0]));
486+
}
487+
int error;
488+
mp_uint_t res = stream_p->ioctl(args[0], MP_STREAM_TRUNCATE, size, &error);
489+
if (res == MP_STREAM_ERROR) {
490+
mp_raise_OSError(error);
491+
}
492+
return mp_obj_new_int(size);
493+
}
494+
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_truncate_obj, 1, 2, stream_truncate);
495+
475496
STATIC mp_obj_t stream_flush(mp_obj_t self) {
476497
const mp_stream_p_t *stream_p = mp_get_stream(self);
477498
int error;

‎py/stream.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
#define MP_STREAM_GET_DATA_OPTS (8) // Get data/message options
4444
#define MP_STREAM_SET_DATA_OPTS (9) // Set data/message options
4545
#define MP_STREAM_GET_FILENO (10) // Get fileno of underlying file
46+
#define MP_STREAM_TRUNCATE (11)
4647

4748
// These poll ioctl values are compatible with Linux
4849
#define MP_STREAM_POLL_RD (0x0001)
@@ -85,6 +86,7 @@ MP_DECLARE_CONST_FUN_OBJ_2(mp_stream_write1_obj);
8586
MP_DECLARE_CONST_FUN_OBJ_1(mp_stream_close_obj);
8687
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_seek_obj);
8788
MP_DECLARE_CONST_FUN_OBJ_1(mp_stream_tell_obj);
89+
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_truncate_obj);
8890
MP_DECLARE_CONST_FUN_OBJ_1(mp_stream_flush_obj);
8991
MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_stream_ioctl_obj);
9092

‎tests/io/file_truncate.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import io
2+
3+
4+
def test_truncate(file):
5+
file.write("a" * 64)
6+
7+
# arg < 0
8+
try:
9+
file.truncate(-1)
10+
except Exception as e:
11+
# Usually this is OSError, but it is ValueError for
12+
# a StringIO in CPython.
13+
print("-1 invalid")
14+
15+
# no arg
16+
file.seek(56)
17+
print(file.truncate()) # 56
18+
print(file.seek(0, 2)) # 56
19+
20+
# position < arg < size: file gets shorter
21+
file.seek(40)
22+
print(file.truncate(48)) # 48
23+
print(file.seek(0, 2)) # 48
24+
25+
# arg < position < size: file gets shorter, position decreases
26+
file.seek(40)
27+
print(file.truncate(32)) # 32
28+
print(file.seek(0, 2)) # 32
29+
print(file.tell()) # 32
30+
31+
# position < size < arg: file gets longer, unless it is StringIO
32+
file.seek(24)
33+
print(file.truncate(64)) # 64
34+
print(file.seek(0, 2)) # 64 or 32
35+
36+
37+
with open("f", "w+") as file:
38+
test_truncate(file)
39+
40+
test_truncate(io.StringIO())

0 commit comments

Comments
 (0)
0