Browse Source

fix base64

blueloveTH 11 tháng trước cách đây
mục cha
commit
95d59cd4c6
3 tập tin đã thay đổi với 14 bổ sung4 xóa
  1. 1 1
      docs/modules/base64.md
  2. 10 2
      src/modules/base64.c
  3. 3 1
      tests/70_base64.py

+ 1 - 1
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`.

+ 10 - 2
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);

+ 3 - 1
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
+assert decoded == data
+
+assert base64.b64decode('8J+llQ==') == '🥕'.encode()