WIP, prototype: Allow to override sys.stdout.

Author: pfalcon

py/modbuiltins.c

@@ -376,7 +376,7 @@ STATIC mp_obj_t mp_builtin_print(mp_uint_t n_args, const mp_obj_t *args, mp_map_
         end_data = mp_obj_str_get_data(end_elem->value, &end_len);
     }
     #if MICROPY_PY_IO
-    mp_obj_t stream_obj = &mp_sys_stdout_obj;
+    mp_obj_t stream_obj = mp_sys_stdout_print.data;
     mp_map_elem_t *file_elem = mp_map_lookup(kwargs, MP_OBJ_NEW_QSTR(MP_QSTR_file), MP_MAP_LOOKUP);
     if (file_elem != NULL && file_elem->value != mp_const_none) {
         stream_obj = file_elem->value;

py/modsys.c

@@ -45,8 +45,12 @@ extern struct _mp_dummy_t mp_sys_stdout_obj;
 extern struct _mp_dummy_t mp_sys_stderr_obj;
 
 #if MICROPY_PY_IO
+#if 1
+mp_print_t mp_sys_stdout_print = {&mp_sys_stdout_obj, (mp_print_strn_t)mp_stream_write};
+#else
 const mp_print_t mp_sys_stdout_print = {&mp_sys_stdout_obj, (mp_print_strn_t)mp_stream_write};
 #endif
+#endif
 
 /// \constant version - Python language version that this implementation conforms to, as a string
 STATIC const MP_DEFINE_STR_OBJ(version_obj, "3.4.0");
@@ -106,7 +110,7 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_sys_exit_obj, 0, 1, mp_sys_exit);
 
 STATIC mp_obj_t mp_sys_print_exception(mp_uint_t n_args, const mp_obj_t *args) {
     #if MICROPY_PY_IO
-    mp_obj_t stream_obj = &mp_sys_stdout_obj;
+    mp_obj_t stream_obj = mp_sys_stdout_print.data;
     if (n_args > 1) {
         stream_obj = args[1];
     }
@@ -197,10 +201,25 @@ STATIC const mp_map_elem_t mp_module_sys_globals_table[] = {
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_sys_globals, mp_module_sys_globals_table);
 
+mp_obj_dict_t *mp_module_sys_override;
+
+STATIC void sys_setattr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
+    if (attr == MP_QSTR_stdout) {
+        mp_sys_stdout_print.data = dest[1];
+        dest[0] = MP_OBJ_NULL;
+    }
+}
+
+STATIC const mp_obj_module_override_helper_t sys_overrides = {
+    .overrides = &mp_module_sys_override,
+    .setattr = sys_setattr,
+};
+
 const mp_obj_module_t mp_module_sys = {
     .base = { &mp_type_module },
     .name = MP_QSTR_sys,
     .globals = (mp_obj_dict_t*)&mp_module_sys_globals,
+    .override_helper = &sys_overrides,
 };
 
 #endif

py/mpprint.h

@@ -51,7 +51,7 @@ typedef struct _mp_print_t {
 extern const mp_print_t mp_plat_print;
 #if MICROPY_PY_IO
 // 2) Wrapper for printing to sys.stdout.
-extern const mp_print_t mp_sys_stdout_print;
+extern /*const*/ mp_print_t mp_sys_stdout_print;
 #endif
 
 int mp_print_str(const mp_print_t *print, const char *str);

py/obj.h

@@ -676,11 +676,19 @@ qstr mp_obj_code_get_name(const byte *code_info);
 mp_obj_t mp_identity(mp_obj_t self);
 MP_DECLARE_CONST_FUN_OBJ(mp_identity_obj);
 
+// Helper structure to support overriding of attributes
+// in a builtin (ROM) module.
+typedef struct _mp_obj_module_override_helper_t {
+    mp_obj_dict_t **overrides;
+    mp_attr_fun_t setattr;
+} mp_obj_module_override_helper_t;
+
 // module
 typedef struct _mp_obj_module_t {
     mp_obj_base_t base;
     qstr name;
     mp_obj_dict_t *globals;
+    const mp_obj_module_override_helper_t *override_helper;
 } mp_obj_module_t;
 mp_obj_dict_t *mp_obj_module_get_globals(mp_obj_t self_in);
 // check if given module object is a package

py/objmodule.c

@@ -33,6 +33,8 @@
 #include "py/runtime.h"
 #include "py/builtin.h"
 
+STATIC const mp_map_t mp_builtin_module_map;
+
 STATIC void module_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
     (void)kind;
     mp_obj_module_t *self = self_in;
@@ -54,6 +56,16 @@ STATIC void module_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     mp_obj_module_t *self = self_in;
     if (dest[0] == MP_OBJ_NULL) {
         // load attribute
+        if (self->override_helper) {
+            mp_obj_dict_t *overrides = *self->override_helper->overrides;
+            if (overrides != MP_OBJ_NULL) {
+                mp_map_elem_t *elem = mp_map_lookup(&overrides->map, MP_OBJ_NEW_QSTR(attr), MP_MAP_LOOKUP);
+                if (elem != NULL) {
+                    dest[0] = elem->value;
+                    return;
+                }
+            }
+        }
         mp_map_elem_t *elem = mp_map_lookup(&self->globals->map, MP_OBJ_NEW_QSTR(attr), MP_MAP_LOOKUP);
         if (elem != NULL) {
             dest[0] = elem->value;
@@ -71,8 +83,23 @@ STATIC void module_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             } else
             #endif
             {
-                // can't delete or store to fixed map
-                return;
+                #if 1
+                if (self->override_helper) {
+                    if (*self->override_helper->overrides != MP_OBJ_NULL) {
+//printf("Found existing override\n");
+                        dict = *self->override_helper->overrides;
+                    } else {
+//printf("Creating new override\n");
+                        dict = mp_obj_new_dict(1);
+                        *self->override_helper->overrides = dict;
+                    }
+                    self->override_helper->setattr(self_in, attr, dest);
+                } else
+                #endif
+                {
+                    // can't delete or store to fixed map
+                    return;
+                }
             }
         }
         if (dest[1] == MP_OBJ_NULL) {
@@ -107,6 +134,7 @@ mp_obj_t mp_obj_new_module(qstr module_name) {
     o->base.type = &mp_type_module;
     o->name = module_name;
     o->globals = mp_obj_new_dict(MICROPY_MODULE_DICT_SIZE);
+    o->override_helper = NULL;
 
     // store __name__ entry in the module
     mp_obj_dict_store(o->globals, MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(module_name));

tests/basics/module2.py

@@ -1,6 +1,6 @@
 # uPy behaviour only: builtin modules are read-only
-import sys
+import micropython
 try:
-    sys.x = 1
+    micropython.x = 1
 except AttributeError:
     print("AttributeError")