#include <stdio.h>

#include <string.h>

#include <assert.h>

#include "py/nlr.h"

#include "py/compile.h"

#include "py/objmodule.h"

#include "py/runtime.h"

#include "py/builtin.h"

#include "py/frozenmod.h"

#if 0 // print debugging info

#define DEBUG_PRINT (1)

#define DEBUG_printf DEBUG_printf

#else // don't print debugging info

#define DEBUG_PRINT (0)

#define DEBUG_printf(...) (void)0

#define PATH_SEP_CHAR '/'

#if MICROPY_MODULE_WEAK_LINKS

STATIC const mp_rom_map_elem_t mp_builtin_module_weak_links_table[] = {

MICROPY_PORT_BUILTIN_MODULE_WEAK_LINKS

};

STATIC MP_DEFINE_CONST_MAP(mp_builtin_module_weak_links_map, mp_builtin_module_weak_links_table);

bool mp_obj_is_package(mp_obj_t module) {

mp_obj_t dest[2];

mp_load_method_maybe(module, MP_QSTR___path__, dest);

return dest[0] != MP_OBJ_NULL;

}

STATIC mp_import_stat_t stat_dir_or_file(vstr_t *path) {

mp_import_stat_t stat = mp_import_stat(vstr_null_terminated_str(path));

DEBUG_printf("stat %s: %d\n", vstr_str(path), stat);

if (stat == MP_IMPORT_STAT_DIR) {

return stat;

}

vstr_add_str(path, ".py");

stat = mp_import_stat(vstr_null_terminated_str(path));

if (stat == MP_IMPORT_STAT_FILE) {

return stat;

}

#if MICROPY_PERSISTENT_CODE_LOAD

vstr_ins_byte(path, path->len - 2, 'm');

stat = mp_import_stat(vstr_null_terminated_str(path));

if (stat == MP_IMPORT_STAT_FILE) {

return stat;

}

#endif

return MP_IMPORT_STAT_NO_EXIST;

}

STATIC mp_import_stat_t find_file(const char *file_str, uint file_len, vstr_t *dest) {

#if MICROPY_PY_SYS

// extract the list of paths

mp_uint_t path_num;

mp_obj_t *path_items;

mp_obj_list_get(mp_sys_path, &path_num, &path_items);

if (path_num == 0) {

// mp_sys_path is empty, so just use the given file name

vstr_add_strn(dest, file_str, file_len);

return stat_dir_or_file(dest);

#if MICROPY_PY_SYS

} else {

// go through each path looking for a directory or file

for (mp_uint_t i = 0; i < path_num; i++) {

vstr_reset(dest);

mp_uint_t p_len;

const char *p = mp_obj_str_get_data(path_items[i], &p_len);

if (p_len > 0) {

vstr_add_strn(dest, p, p_len);

vstr_add_char(dest, PATH_SEP_CHAR);

}

vstr_add_strn(dest, file_str, file_len);

mp_import_stat_t stat = stat_dir_or_file(dest);

if (stat != MP_IMPORT_STAT_NO_EXIST) {

return stat;

}

}

// could not find a directory or file

return MP_IMPORT_STAT_NO_EXIST;

}

}

#if MICROPY_ENABLE_COMPILER

STATIC void do_load_from_lexer(mp_obj_t module_obj, mp_lexer_t *lex, const char *fname) {

if (lex == NULL) {

// we verified the file exists using stat, but lexer could still fail

if (MICROPY_ERROR_REPORTING == MICROPY_ERROR_REPORTING_TERSE) {

nlr_raise(mp_obj_new_exception_msg(&mp_type_ImportError, "module not found"));

} else {

nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_ImportError,

"no module named '%s'", fname));

}

}

#if MICROPY_PY___FILE__

qstr source_name = lex->source_name;

mp_store_attr(module_obj, MP_QSTR___file__, MP_OBJ_NEW_QSTR(source_name));

#endif

// parse, compile and execute the module in its context

mp_obj_dict_t *mod_globals = mp_obj_module_get_globals(module_obj);

mp_parse_compile_execute(lex, MP_PARSE_FILE_INPUT, mod_globals, mod_globals);

}

#if MICROPY_PERSISTENT_CODE_LOAD

STATIC void do_execute_raw_code(mp_obj_t module_obj, mp_raw_code_t *raw_code) {

#if MICROPY_PY___FILE__

// TODO

//qstr source_name = lex->source_name;

//mp_store_attr(module_obj, MP_QSTR___file__, MP_OBJ_NEW_QSTR(source_name));

#endif

// execute the module in its context

mp_obj_dict_t *mod_globals = mp_obj_module_get_globals(module_obj);

// save context

mp_obj_dict_t *volatile old_globals = mp_globals_get();

mp_obj_dict_t *volatile old_locals = mp_locals_get();

// set new context

mp_globals_set(mod_globals);

mp_locals_set(mod_globals);

nlr_buf_t nlr;

if (nlr_push(&nlr) == 0) {

mp_obj_t module_fun = mp_make_function_from_raw_code(raw_code, MP_OBJ_NULL, MP_OBJ_NULL);

mp_call_function_0(module_fun);

// finish nlr block, restore context

nlr_pop();

mp_globals_set(old_globals);

mp_locals_set(old_locals);

} else {

// exception; restore context and re-raise same exception

mp_globals_set(old_globals);

mp_locals_set(old_locals);

nlr_jump(nlr.ret_val);

}

}

STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {

// create the lexer

char *file_str = vstr_null_terminated_str(file);

#if MICROPY_PERSISTENT_CODE_LOAD

if (file_str[file->len - 3] == 'm') {

mp_raw_code_t *raw_code = mp_raw_code_load_file(file_str);

do_execute_raw_code(module_obj, raw_code);

return;

}

#endif

#if MICROPY_ENABLE_COMPILER

{

mp_lexer_t *lex = mp_lexer_new_from_file(file_str);

do_load_from_lexer(module_obj, lex, file_str);

}

#else

nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_ImportError,

"script compilation not supported"));

#endif

}

STATIC void chop_component(const char *start, const char **end) {

const char *p = *end;

while (p > start) {

if (*--p == '.') {

*end = p;

return;

}

}

*end = p;

}

mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {

#if DEBUG_PRINT

DEBUG_printf("__import__:\n");

for (mp_uint_t i = 0; i < n_args; i++) {

DEBUG_printf(" ");

mp_obj_print(args[i], PRINT_REPR);

DEBUG_printf("\n");

}

mp_obj_t module_name = args[0];

mp_obj_t fromtuple = mp_const_none;

mp_int_t level = 0;

if (n_args >= 4) {

fromtuple = args[3];

if (n_args >= 5) {

level = MP_OBJ_SMALL_INT_VALUE(args[4]);

}

}

mp_uint_t mod_len;

const char *mod_str = mp_obj_str_get_data(module_name, &mod_len);

if (level != 0) {

// What we want to do here is to take name of current module,

// chop <level> trailing components, and concatenate with passed-in

// module name, thus resolving relative import name into absolue.

// This even appears to be correct per

// http://legacy.python.org/dev/peps/pep-0328/#relative-imports-and-name

// "Relative imports use a module's __name__ attribute to determine that

// module's position in the package hierarchy."

level--;

mp_obj_t this_name_q = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___name__));

assert(this_name_q != MP_OBJ_NULL);

#if MICROPY_CPYTHON_COMPAT

if (MP_OBJ_QSTR_VALUE(this_name_q) == MP_QSTR___main__) {

// This is a module run by -m command-line switch, get its real name from backup attribute

this_name_q = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___main__));

}

#endif

mp_map_t *globals_map = &mp_globals_get()->map;

mp_map_elem_t *elem = mp_map_lookup(globals_map, MP_OBJ_NEW_QSTR(MP_QSTR___path__), MP_MAP_LOOKUP);

bool is_pkg = (elem != NULL);

#if DEBUG_PRINT

DEBUG_printf("Current module/package: ");

mp_obj_print(this_name_q, PRINT_REPR);

DEBUG_printf(", is_package: %d", is_pkg);

DEBUG_printf("\n");

mp_uint_t this_name_l;

const char *this_name = mp_obj_str_get_data(this_name_q, &this_name_l);

const char *p = this_name + this_name_l;

if (!is_pkg) {

// We have module, but relative imports are anchored at package, so

// go there.

chop_component(this_name, &p);

}

uint dots_seen = 0;

while (level--) {

chop_component(this_name, &p);

dots_seen++;

}

if (dots_seen == 0 && level >= 1) {

// http://legacy.python.org/dev/peps/pep-0328/#relative-imports-and-name

// "If the module's name does not contain any package information

// (e.g. it is set to '__main__') then relative imports are

// resolved as if the module were a top level module, regardless

// of where the module is actually located on the file system."

// Supposedly this if catches this condition and resolve it properly

// TODO: But nobody knows for sure. This condition happens when

// package's __init__.py does something like "import .submod". So,

// maybe we should check for package here? But quote above doesn't

// talk about packages, it talks about dot-less module names.

DEBUG_printf("Warning: no dots in current module name and level>0\n");

p = this_name + this_name_l;

} else if (level != -1) {

nlr_raise(mp_obj_new_exception_msg(&mp_type_ImportError, "Invalid relative import"));

}

uint new_mod_l = (mod_len == 0 ? (size_t)(p - this_name) : (size_t)(p - this_name) + 1 + mod_len);

char *new_mod = alloca(new_mod_l);

memcpy(new_mod, this_name, p - this_name);

if (mod_len != 0) {

new_mod[p - this_name] = '.';

memcpy(new_mod + (p - this_name) + 1, mod_str, mod_len);

}

qstr new_mod_q = qstr_from_strn(new_mod, new_mod_l);

DEBUG_printf("Resolved base name for relative import: '%s'\n", qstr_str(new_mod_q));

if (new_mod_q == MP_QSTR_) {

// CPython raises SystemError

nlr_raise(mp_obj_new_exception_msg(&mp_type_ImportError, "cannot perform relative import"));

}

module_name = MP_OBJ_NEW_QSTR(new_mod_q);

mod_str = new_mod;

mod_len = new_mod_l;

}

// check if module already exists

qstr module_name_qstr = mp_obj_str_get_qstr(module_name);

mp_obj_t module_obj = mp_module_get(module_name_qstr);

if (module_obj != MP_OBJ_NULL) {

DEBUG_printf("Module already loaded\n");

// If it's not a package, return module right away

char *p = strchr(mod_str, '.');

if (p == NULL) {

return module_obj;

}

// If fromlist is not empty, return leaf module

if (fromtuple != mp_const_none) {

return module_obj;

}

// Otherwise, we need to return top-level package

qstr pkg_name = qstr_from_strn(mod_str, p - mod_str);

return mp_module_get(pkg_name);

}

DEBUG_printf("Module not yet loaded\n");

#if MICROPY_MODULE_FROZEN

mp_lexer_t *lex = mp_find_frozen_module(mod_str, mod_len);

if (lex != NULL) {

module_obj = mp_obj_new_module(module_name_qstr);

// if args[3] (fromtuple) has magic value False, set up

// this module for command-line "-m" option (set module's

// name to __main__ instead of real name).

// TODO: Duplicated below too.

if (fromtuple == mp_const_false) {

mp_obj_module_t *o = MP_OBJ_TO_PTR(module_obj);

mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR___main__));

}

do_load_from_lexer(module_obj, lex, mod_str);

return module_obj;

}

#endif

uint last = 0;

VSTR_FIXED(path, MICROPY_ALLOC_PATH_MAX)

module_obj = MP_OBJ_NULL;

mp_obj_t top_module_obj = MP_OBJ_NULL;

mp_obj_t outer_module_obj = MP_OBJ_NULL;

uint i;

for (i = 1; i <= mod_len; i++) {

if (i == mod_len || mod_str[i] == '.') {

// create a qstr for the module name up to this depth

qstr mod_name = qstr_from_strn(mod_str, i);

DEBUG_printf("Processing module: %s\n", qstr_str(mod_name));

DEBUG_printf("Previous path: =%.*s=\n", vstr_len(&path), vstr_str(&path));

// find the file corresponding to the module name

mp_import_stat_t stat;

if (vstr_len(&path) == 0) {

// first module in the dotted-name; search for a directory or file

stat = find_file(mod_str, i, &path);

} else {

// latter module in the dotted-name; append to path

vstr_add_char(&path, PATH_SEP_CHAR);

vstr_add_strn(&path, mod_str + last, i - last);

stat = stat_dir_or_file(&path);

}

DEBUG_printf("Current path: %.*s\n", vstr_len(&path), vstr_str(&path));

if (stat == MP_IMPORT_STAT_NO_EXIST) {

#if MICROPY_MODULE_WEAK_LINKS

// check if there is a weak link to this module

if (i == mod_len) {

mp_map_elem_t *el = mp_map_lookup((mp_map_t*)&mp_builtin_module_weak_links_map, MP_OBJ_NEW_QSTR(mod_name), MP_MAP_LOOKUP);

if (el == NULL) {

goto no_exist;

}

// found weak linked module

module_obj = el->value;

} else {

no_exist:

#else

{

#endif

// couldn't find the file, so fail

if (MICROPY_ERROR_REPORTING == MICROPY_ERROR_REPORTING_TERSE) {

nlr_raise(mp_obj_new_exception_msg(&mp_type_ImportError, "module not found"));

} else {

nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_ImportError,

"no module named '%q'", mod_name));

}

}

} else {

// found the file, so get the module

module_obj = mp_module_get(mod_name);

}

if (module_obj == MP_OBJ_NULL) {

// module not already loaded, so load it!

module_obj = mp_obj_new_module(mod_name);

// if args[3] (fromtuple) has magic value False, set up

// this module for command-line "-m" option (set module's

// name to __main__ instead of real name).

if (i == mod_len && fromtuple == mp_const_false) {

mp_obj_module_t *o = MP_OBJ_TO_PTR(module_obj);

mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR___main__));

#if MICROPY_CPYTHON_COMPAT

// Store real name in "__main__" attribute. Choosen semi-randonly, to reuse existing qstr's.

mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___main__), MP_OBJ_NEW_QSTR(mod_name));

#endif

}

if (stat == MP_IMPORT_STAT_DIR) {

DEBUG_printf("%.*s is dir\n", vstr_len(&path), vstr_str(&path));

// https://docs.python.org/3/reference/import.html

// "Specifically, any module that contains a __path__ attribute is considered a package."

mp_store_attr(module_obj, MP_QSTR___path__, mp_obj_new_str(vstr_str(&path), vstr_len(&path), false));

vstr_add_char(&path, PATH_SEP_CHAR);

vstr_add_str(&path, "__init__.py");

if (mp_import_stat(vstr_null_terminated_str(&path)) != MP_IMPORT_STAT_FILE) {

vstr_cut_tail_bytes(&path, sizeof("/__init__.py") - 1); // cut off /__init__.py

mp_warning("%s is imported as namespace package", vstr_str(&path));

} else {

do_load(module_obj, &path);

vstr_cut_tail_bytes(&path, sizeof("/__init__.py") - 1); // cut off /__init__.py

}

} else { // MP_IMPORT_STAT_FILE

do_load(module_obj, &path);

// TODO: We cannot just break here, at the very least, we must execute

// trailer code below. But otherwise if there're remaining components,

// that would be (??) object path within module, not modules path within FS.

// break;

}

}

if (outer_module_obj != MP_OBJ_NULL) {

qstr s = qstr_from_strn(mod_str + last, i - last);

mp_store_attr(outer_module_obj, s, module_obj);

}

outer_module_obj = module_obj;

if (top_module_obj == MP_OBJ_NULL) {

top_module_obj = module_obj;

}

last = i + 1;

}

}

if (i < mod_len) {

// we loaded a package, now need to load objects from within that package

// TODO

assert(0);

}

// If fromlist is not empty, return leaf module

if (fromtuple != mp_const_none) {

return module_obj;

}

// Otherwise, we need to return top-level package

return top_module_obj;

}

MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin___import___obj, 1, 5, mp_builtin___import__);