diff --git a/py/compile.c b/py/compile.c index 62757de3c0..b2489a5f1a 100644 --- a/py/compile.c +++ b/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; } diff --git a/py/objfun.c b/py/objfun.c index 1ebfa3d5af..9b53f42306 100644 --- a/py/objfun.c +++ b/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)); } diff --git a/py/objfun.h b/py/objfun.h index af7c334858..a4f035d5af 100644 --- a/py/objfun.h +++ b/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