From 95d59cd4c6e787366f9f411fd572039dbbb8bf98 Mon Sep 17 00:00:00 2001 From: blueloveTH Date: Sat, 5 Apr 2025 03:07:24 +0800 Subject: [PATCH] fix base64 --- docs/modules/base64.md | 2 +- src/modules/base64.c | 12 ++++++++++-- tests/70_base64.py | 4 +++- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/modules/base64.md b/docs/modules/base64.md index 8d186b17..25d0b969 100644 --- a/docs/modules/base64.md +++ b/docs/modules/base64.md @@ -7,6 +7,6 @@ label: base64 Encode bytes-like object `b` using the standard Base64 alphabet. -### `base64.b64decode(b: bytes) -> bytes` +### `base64.b64decode(b: str | bytes) -> bytes` Decode Base64 encoded bytes-like object `b`. diff --git a/src/modules/base64.c b/src/modules/base64.c index bda91126..5a51b42c 100644 --- a/src/modules/base64.c +++ b/src/modules/base64.c @@ -178,9 +178,17 @@ static bool base64_b64encode(int argc, py_Ref argv) { static bool base64_b64decode(int argc, py_Ref argv) { PY_CHECK_ARGC(1); - PY_CHECK_ARG_TYPE(0, tp_bytes); int src_size; - unsigned char* src_data = py_tobytes(argv, &src_size); + void* src_data; + if(py_istype(argv, tp_str)) { + c11_sv sv = py_tosv(argv); + src_data = (void*)sv.data; + src_size = sv.size; + } else if(py_istype(argv, tp_bytes)) { + src_data = py_tobytes(argv, &src_size); + } else { + return TypeError("expect bytes or str, got %t", argv->type); + } unsigned char* dst_data = py_newbytes(py_retval(), src_size); int size = base64_decode((const char*)src_data, src_size, dst_data); py_bytes_resize(py_retval(), size); diff --git a/tests/70_base64.py b/tests/70_base64.py index 6ca029c9..0cb9df92 100644 --- a/tests/70_base64.py +++ b/tests/70_base64.py @@ -17,4 +17,6 @@ res = 'Qm6RErAN/8qtbbLCq8aPGHEuRl5HjJ+/hua+4N9e2RTximh4+VuGMGwxAPnr4eS+P8zYZpkzT assert encoded.decode() == res decoded = base64.b64decode(encoded) -assert decoded == data \ No newline at end of file +assert decoded == data + +assert base64.b64decode('8J+llQ==') == '🥕'.encode()