Compile simple python packages and load it into the ASR

Author: Thirumalai-Shaktivel

src/lpython/semantics/python_ast_to_asr.cpp

@@ -307,6 +307,7 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
     bool compile_module = true;
     ASR::TranslationUnit_t* mod1 = nullptr;
     std::string input;
+    std::string module_dir_name = "";
     bool found = set_module_path(infile0c, rl_path, infile,
                                  path_used, input, ltypes, enum_py);
     if( !found ) {
@@ -335,7 +336,15 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
 
     // insert into `symtab`
     std::vector<std::pair<std::string, ASR::Module_t*>> children_modules;
-    ASRUtils::extract_module_python(*mod1, children_modules, module_name);
+    if (module_name == "__init__") {
+        // remove `__init__.py`
+        module_dir_name = infile.substr(0, infile.find_last_of('/'));
+        // assign module directory name
+        module_dir_name = module_dir_name.substr(module_dir_name.find_last_of('/') + 1);
+        ASRUtils::extract_module_python(*mod1, children_modules, module_dir_name);
+    } else {
+        ASRUtils::extract_module_python(*mod1, children_modules, module_name);
+    }
     ASR::Module_t* mod2 = nullptr;
     for( auto& a: children_modules ) {
         std::string a_name = a.first;
@@ -344,6 +353,10 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
             a_mod->m_name = s2c(al, module_name);
             a_mod->m_intrinsic = intrinsic;
             mod2 = a_mod;
+        } else if (a_name == module_dir_name) {
+            a_mod->m_name = s2c(al, module_dir_name);
+            a_mod->m_intrinsic = intrinsic;
+            mod2 = a_mod;
         }
         symtab->add_symbol(a_name, (ASR::symbol_t*)a_mod);
         a_mod->m_symtab->parent = symtab;
@@ -429,9 +442,18 @@ ASR::symbol_t* import_from_module(Allocator &al, ASR::Module_t *m, SymbolTable *
             ASR::accessType::Public
             );
         return ASR::down_cast<ASR::symbol_t>(v);
+    } else if (ASR::is_a<ASR::ExternalSymbol_t>(*t)) {
+        ASR::ExternalSymbol_t *es = ASR::down_cast<ASR::ExternalSymbol_t>(t);
+        SymbolTable *symtab = current_scope;
+        while (symtab->parent != nullptr) symtab = symtab->parent;
+        ASR::symbol_t *sym = symtab->get_symbol(es->m_module_name);
+        ASR::Module_t *m = ASR::down_cast<ASR::Module_t>(sym);
+
+        return import_from_module(al, m, symtab, es->m_name,
+                            cur_sym_name, new_sym_name, loc);
     } else {
-        throw SemanticError("Only Subroutines, Functions and Variables are currently supported in 'import'",
-            loc);
+        throw SemanticError("Only Subroutines, Functions, Variables and "
+            "ExternalSymbol are currently supported in 'import'", loc);
     }
     LFORTRAN_ASSERT(false);
     return nullptr;
@@ -3157,6 +3179,24 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
                 st = st->parent;
             }
             bool ltypes, enum_py;
+            if (msym != "ltypes") {
+                if (import_path != "" &&
+                        !path_exits(paths[0] + '/' + msym + ".py")) {
+                    paths = {import_path};
+                    if (parent_dir != "") paths[0] += '/' + parent_dir;
+                    if(is_directory(paths[0])) {
+                        paths[0] += '/' + msym;
+                        msym = "__init__";
+                    }
+                } else if (is_directory(msym)) {
+                    paths = {rl_path, msym};
+                    msym = "__init__";
+                } else if (paths[1] != ""
+                        && is_directory(paths[1] + '/' + msym)) {
+                    paths[1] += '/' + msym;
+                    msym = "__init__";
+                }
+            }
             t = (ASR::symbol_t*)(load_module(al, st,
                 msym, x.base.base.loc, false, paths, ltypes, enum_py,
                 [&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); }
@@ -3202,6 +3242,24 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
         }
         for (auto &mod_sym : mods) {
             bool ltypes, enum_py;
+            if (mod_sym != "ltypes") {
+                if (import_path != "" &&
+                        !path_exits(paths[0] + '/' + mod_sym + ".py")) {
+                    paths = {import_path};
+                    if (parent_dir != "") paths[0] += '/' + parent_dir;
+                    if(is_directory(paths[0])) {
+                        paths[0] += '/' + mod_sym;
+                        mod_sym = "__init__";
+                    }
+                } else if (is_directory(mod_sym)) {
+                    paths = {rl_path, mod_sym};
+                    mod_sym = "__init__";
+                } else if (paths[1] != ""
+                        && is_directory(paths[1] + '/' + mod_sym)) {
+                    paths[1] += '/' + mod_sym;
+                    mod_sym = "__init__";
+                }
+            }
             t = (ASR::symbol_t*)(load_module(al, st,
                 mod_sym, x.base.base.loc, false, paths, ltypes, enum_py,
                 [&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); }