Support __bytes

tannewt · tannewt · commit 9435e01f9e46 · 2019-10-14T16:05:17.000-07:00
Fixes adafruit#1763
diff --git a/py/objstr.c b/py/objstr.c
@@ -31,6 +31,7 @@
 #include "py/unicode.h"
 #include "py/objstr.h"
 #include "py/objlist.h"
+#include "py/objtype.h"
 #include "py/runtime.h"
 #include "py/stackctrl.h"
 
@@ -211,6 +212,10 @@ STATIC mp_obj_t bytes_make_new(const mp_obj_type_t *type_in, size_t n_args, cons
         return mp_const_empty_bytes;
     }
 
+    if (n_args > 1) {
+        goto wrong_args;
+    }
+
     if (MP_OBJ_IS_STR(args[0])) {
         if (n_args < 2 || n_args > 3) {
             goto wrong_args;
@@ -226,10 +231,6 @@ STATIC mp_obj_t bytes_make_new(const mp_obj_type_t *type_in, size_t n_args, cons
         return MP_OBJ_FROM_PTR(o);
     }
 
-    if (n_args > 1) {
-        goto wrong_args;
-    }
-
     if (MP_OBJ_IS_SMALL_INT(args[0])) {
         mp_int_t len = MP_OBJ_SMALL_INT_VALUE(args[0]);
         if (len < 0) {
@@ -241,6 +242,13 @@ STATIC mp_obj_t bytes_make_new(const mp_obj_type_t *type_in, size_t n_args, cons
         return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
     }
 
+    // check if __bytes__ exists, and if so delegate to it
+    mp_obj_t dest[2];
+    mp_load_method_maybe(args[0], MP_QSTR___bytes__, dest);
+    if (dest[0] != MP_OBJ_NULL) {
+        return mp_call_method_n_kw(0, 0, dest);
+    }
+
     // check if argument has the buffer protocol
     mp_buffer_info_t bufinfo;
     if (mp_get_buffer(args[0], &bufinfo, MP_BUFFER_READ)) {
diff --git a/tests/basics/class_bytes.py b/tests/basics/class_bytes.py
@@ -0,0 +1,9 @@
+class C1:
+    def __init__(self, value):
+        self.value = value
+
+    def __bytes__(self):
+        return self.value
+
+c1 = C1(b"class 1")
+print(bytes(c1))
