Merge 0a3f167fdf into 5114f2c1ea

py/compile.c

@@ -786,6 +786,32 @@ static qstr compile_funcdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns
     // compile the function definition
     compile_funcdef_lambdef(comp, fscope, pns->nodes[1], PN_typedargslist);
 
+    #if MICROPY_ENABLE_DOC_STRING
+    // look for the first statement
+    mp_parse_node_t pn = pns->nodes[3]; // the function body/suite
+    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_suite_block_stmts)) {
+        // a list of statements; get the first one
+        pn = ((mp_parse_node_struct_t *)pn)->nodes[0];
+    }
+
+    // check the first statement for a doc string
+    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_expr_stmt)) {
+        mp_parse_node_struct_t *pns2 = (mp_parse_node_struct_t *)pn;
+        if ((MP_PARSE_NODE_IS_LEAF(pns2->nodes[0])
+             && MP_PARSE_NODE_LEAF_KIND(pns2->nodes[0]) == MP_PARSE_NODE_STRING)
+            || (MP_PARSE_NODE_IS_STRUCT_KIND(pns2->nodes[0], PN_const_object)
+                && mp_obj_is_str(mp_parse_node_extract_const_object((mp_parse_node_struct_t *)pns2->nodes[0])))) {
+            // duplicate the function on the stack
+            EMIT(dup_top);
+            // compile the doc string
+            compile_node(comp, pns2->nodes[0]);
+            // store the doc string into the function
+            EMIT(rot_two);
+            EMIT_ARG(attr, MP_QSTR___doc__, MP_EMIT_ATTR_STORE);
+        }
+    }
+    #endif
+
     // return its name (the 'f' in "def f(...):")
     return fscope->simple_name;
 }

py/objfun.c

@@ -329,6 +329,17 @@ static mp_obj_t fun_bc_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const
 void mp_obj_fun_bc_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     if (dest[0] != MP_OBJ_NULL) {
         // not load attribute
+        #if MICROPY_ENABLE_DOC_STRING
+        if (attr == MP_QSTR___doc__) {
+            mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
+            if (dest[1] == MP_OBJ_NULL) {
+                self->doc_string = mp_const_none;
+            } else {
+                self->doc_string = dest[1];
+            }
+            dest[0] = MP_OBJ_NULL; // success with store/delete
+        }
+        #endif
         return;
     }
     if (attr == MP_QSTR___name__) {
@@ -338,6 +349,12 @@ void mp_obj_fun_bc_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
         mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
         dest[0] = MP_OBJ_FROM_PTR(self->context->module.globals);
     }
+    #if MICROPY_ENABLE_DOC_STRING
+    if (attr == MP_QSTR___doc__) {
+        mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
+        dest[0] = self->doc_string;
+    }
+    #endif
 }
 #endif
 
@@ -382,6 +399,9 @@ mp_obj_t mp_obj_new_fun_bc(const mp_obj_t *def_args, const byte *code, const mp_
     o->bytecode = code;
     o->context = context;
     o->child_table = child_table;
+    #if MICROPY_ENABLE_DOC_STRING
+    o->doc_string = mp_const_none;
+    #endif
     if (def_pos_args != NULL) {
         memcpy(o->extra_args, def_pos_args->items, n_def_args * sizeof(mp_obj_t));
     }

py/objfun.h

@@ -34,6 +34,9 @@ typedef struct _mp_obj_fun_bc_t {
     const mp_module_context_t *context;         // context within which this function was defined
     struct _mp_raw_code_t *const *child_table;  // table of children
     const byte *bytecode;                       // bytecode for the function
+    #if MICROPY_ENABLE_DOC_STRING
+    mp_obj_t doc_string;
+    #endif
     #if MICROPY_PY_SYS_SETTRACE
     const struct _mp_raw_code_t *rc;
     #endif