Skip to content

Commit

Permalink
Allow importing python modules from same directory directly
Browse files Browse the repository at this point in the history
  • Loading branch information
Smit-create committed Jun 19, 2022
1 parent bd8b718 commit 999a189
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 27 deletions.
10 changes: 5 additions & 5 deletions src/bin/lpython.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ int emit_asr(const std::string &infile,
diagnostics.diagnostics.clear();
LFortran::Result<LFortran::ASR::TranslationUnit_t*>
r = LFortran::LPython::python_ast_to_asr(al, *ast, diagnostics, true,
compiler_options.symtab_only);
compiler_options.symtab_only, infile);
std::cerr << diagnostics.render(input, lm, compiler_options);
if (!r.ok) {
LFORTRAN_ASSERT(diagnostics.has_error())
Expand Down Expand Up @@ -205,7 +205,7 @@ int emit_cpp(const std::string &infile,
diagnostics.diagnostics.clear();
LFortran::Result<LFortran::ASR::TranslationUnit_t*>
r1 = LFortran::LPython::python_ast_to_asr(al, *ast, diagnostics, true,
compiler_options.symtab_only);
compiler_options.symtab_only, infile);
std::cerr << diagnostics.render(input, lm, compiler_options);
if (!r1.ok) {
LFORTRAN_ASSERT(diagnostics.has_error())
Expand Down Expand Up @@ -245,7 +245,7 @@ int emit_c(const std::string &infile,
diagnostics.diagnostics.clear();
LFortran::Result<LFortran::ASR::TranslationUnit_t*>
r1 = LFortran::LPython::python_ast_to_asr(al, *ast, diagnostics, true,
compiler_options.symtab_only);
compiler_options.symtab_only, infile);
std::cerr << diagnostics.render(input, lm, compiler_options);
if (!r1.ok) {
LFORTRAN_ASSERT(diagnostics.has_error())
Expand Down Expand Up @@ -288,7 +288,7 @@ int emit_llvm(const std::string &infile,
diagnostics.diagnostics.clear();
LFortran::Result<LFortran::ASR::TranslationUnit_t*>
r1 = LFortran::LPython::python_ast_to_asr(al, *ast, diagnostics, true,
compiler_options.symtab_only);
compiler_options.symtab_only, infile);
std::cerr << diagnostics.render(input, lm, compiler_options);
if (!r1.ok) {
LFORTRAN_ASSERT(diagnostics.has_error())
Expand Down Expand Up @@ -352,7 +352,7 @@ int compile_python_to_object_file(
auto ast_to_asr_start = std::chrono::high_resolution_clock::now();
LFortran::Result<LFortran::ASR::TranslationUnit_t*>
r1 = LFortran::LPython::python_ast_to_asr(al, *ast, diagnostics, true,
compiler_options.symtab_only);
compiler_options.symtab_only, infile);
auto ast_to_asr_end = std::chrono::high_resolution_clock::now();
times.push_back(std::make_pair("AST to ASR", std::chrono::duration<double, std::milli>(ast_to_asr_end - ast_to_asr_start).count()));
std::cerr << diagnostics.render(input, lm, compiler_options);
Expand Down
54 changes: 34 additions & 20 deletions src/lpython/semantics/python_ast_to_asr.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ LFortran::Result<std::string> get_full_path(const std::string &filename,
ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
const std::string &module_name,
const Location &loc, bool intrinsic,
const std::string &rl_path,
std::vector<std::string> &rl_path,
bool &ltypes, bool &numpy,
const std::function<void (const std::string &, const Location &)> err) {
ltypes = false;
Expand All @@ -93,17 +93,26 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,

// Parse the module `module_name`.py to AST
std::string infile0 = module_name + ".py";
Result<std::string> rinfile = get_full_path(infile0, rl_path, ltypes,
numpy);
if (!rinfile.ok) {
bool found = false;
std::string path_used = "", infile;
for (auto path: rl_path) {
Result<std::string> rinfile = get_full_path(infile0, path, ltypes, numpy);
if (rinfile.ok) {
found = true;
infile = rinfile.result;
path_used = path;
break;
}
}
if (!found) {
err("Could not find the module '" + infile0 + "'", loc);
}
if (ltypes) return nullptr;
if (numpy) return nullptr;
std::string infile = rinfile.result;

// TODO: diagnostic should be an argument to this function
diag::Diagnostics diagnostics;
Result<AST::ast_t*> r = parse_python_file(al, rl_path, infile,
Result<AST::ast_t*> r = parse_python_file(al, rl_path[0], infile,
diagnostics, false);
if (!r.ok) {
err("The file '" + infile + "' failed to parse", loc);
Expand All @@ -113,7 +122,8 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
// Convert the module from AST to ASR
LFortran::LocationManager lm;
lm.in_filename = infile;
Result<ASR::TranslationUnit_t*> r2 = python_ast_to_asr(al, *ast, diagnostics, false, false);
Result<ASR::TranslationUnit_t*> r2 = python_ast_to_asr(al, *ast, diagnostics,
false, false, path_used);
std::string input;
read_file(infile, input);
CompilerOptions compiler_options;
Expand Down Expand Up @@ -269,12 +279,13 @@ class CommonVisitor : public AST::BaseVisitor<Derived> {
PythonIntrinsicProcedures intrinsic_procedures;
AttributeHandler attr_handler;
std::map<int, ASR::symbol_t*> &ast_overload;
std::string parent_dir;

CommonVisitor(Allocator &al, SymbolTable *symbol_table,
diag::Diagnostics &diagnostics, bool main_module,
std::map<int, ASR::symbol_t*> &ast_overload)
std::map<int, ASR::symbol_t*> &ast_overload, std::string parent_dir)
: diag{diagnostics}, al{al}, current_scope{symbol_table}, main_module{main_module},
ast_overload{ast_overload} {
ast_overload{ast_overload}, parent_dir{parent_dir} {
current_module_dependencies.reserve(al, 4);
}

Expand All @@ -296,9 +307,10 @@ class CommonVisitor : public AST::BaseVisitor<Derived> {

SymbolTable *tu_symtab = ASRUtils::get_tu_symtab(current_scope);
std::string rl_path = get_runtime_library_dir();
std::vector<std::string> paths = {rl_path, parent_dir};
bool ltypes, numpy;
ASR::Module_t *m = load_module(al, tu_symtab, module_name,
loc, true, rl_path,
loc, true, paths,
ltypes, numpy,
[&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); }
);
Expand Down Expand Up @@ -1905,8 +1917,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {

SymbolTableVisitor(Allocator &al, SymbolTable *symbol_table,
diag::Diagnostics &diagnostics, bool main_module,
std::map<int, ASR::symbol_t*> &ast_overload)
: CommonVisitor(al, symbol_table, diagnostics, main_module, ast_overload), is_derived_type{false} {}
std::map<int, ASR::symbol_t*> &ast_overload, std::string parent_dir)
: CommonVisitor(al, symbol_table, diagnostics, main_module, ast_overload, parent_dir), is_derived_type{false} {}


ASR::symbol_t* resolve_symbol(const Location &loc, const std::string &sub_name) {
Expand Down Expand Up @@ -2130,12 +2142,13 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
if (!t) {
std::string rl_path = get_runtime_library_dir();
SymbolTable *st = current_scope;
std::vector<std::string> paths = {rl_path, parent_dir};
if (!main_module) {
st = st->parent;
}
bool ltypes, numpy;
t = (ASR::symbol_t*)(load_module(al, st,
msym, x.base.base.loc, false, rl_path, ltypes, numpy,
msym, x.base.base.loc, false, paths, ltypes, numpy,
[&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); }
));
if (ltypes || numpy) {
Expand Down Expand Up @@ -2164,6 +2177,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
void visit_Import(const AST::Import_t &x) {
ASR::symbol_t *t = nullptr;
std::string rl_path = get_runtime_library_dir();
std::vector<std::string> paths = {rl_path, parent_dir};
SymbolTable *st = current_scope;
std::vector<std::string> mods;
for (size_t i=0; i<x.n_names; i++) {
Expand All @@ -2175,7 +2189,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
for (auto &mod_sym : mods) {
bool ltypes, numpy;
t = (ASR::symbol_t*)(load_module(al, st,
mod_sym, x.base.base.loc, false, rl_path, ltypes, numpy,
mod_sym, x.base.base.loc, false, paths, ltypes, numpy,
[&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); }
));
if (ltypes || numpy) {
Expand Down Expand Up @@ -2206,9 +2220,9 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {

Result<ASR::asr_t*> symbol_table_visitor(Allocator &al, const AST::Module_t &ast,
diag::Diagnostics &diagnostics, bool main_module,
std::map<int, ASR::symbol_t*> &ast_overload)
std::map<int, ASR::symbol_t*> &ast_overload, std::string parent_dir)
{
SymbolTableVisitor v(al, nullptr, diagnostics, main_module, ast_overload);
SymbolTableVisitor v(al, nullptr, diagnostics, main_module, ast_overload, parent_dir);
try {
v.visit_Module(ast);
} catch (const SemanticError &e) {
Expand All @@ -2232,7 +2246,7 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {

BodyVisitor(Allocator &al, ASR::asr_t *unit, diag::Diagnostics &diagnostics,
bool main_module, std::map<int, ASR::symbol_t*> &ast_overload)
: CommonVisitor(al, nullptr, diagnostics, main_module, ast_overload), asr{unit} {}
: CommonVisitor(al, nullptr, diagnostics, main_module, ast_overload, ""), asr{unit} {}

// Transforms statements to a list of ASR statements
// In addition, it also inserts the following nodes if needed:
Expand Down Expand Up @@ -3567,15 +3581,15 @@ std::string pickle_tree_python(AST::ast_t &ast, bool colors) {

Result<ASR::TranslationUnit_t*> python_ast_to_asr(Allocator &al,
AST::ast_t &ast, diag::Diagnostics &diagnostics, bool main_module,
bool symtab_only)
bool symtab_only, const std::string &file_path)
{
std::map<int, ASR::symbol_t*> ast_overload;

std::filesystem::path fp = file_path;
AST::Module_t *ast_m = AST::down_cast2<AST::Module_t>(&ast);

ASR::asr_t *unit;
auto res = symbol_table_visitor(al, *ast_m, diagnostics, main_module,
ast_overload);
ast_overload, fp.parent_path());
if (res.ok) {
unit = res.result;
} else {
Expand Down
4 changes: 2 additions & 2 deletions src/lpython/semantics/python_ast_to_asr.h
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,10 @@
namespace LFortran::LPython {

std::string pickle_python(AST::ast_t &ast, bool colors=false, bool indent=false);
std::string pickle_tree_python(AST::ast_t &ast, bool colors=true);
std::string pickle_tree_python(AST::ast_t &ast, bool colors=true);
Result<ASR::TranslationUnit_t*> python_ast_to_asr(Allocator &al,
LPython::AST::ast_t &ast, diag::Diagnostics &diagnostics, bool main_module,
bool symtab_only);
bool symtab_only, const std::string &file_path);

} // namespace LFortran

Expand Down

0 comments on commit 999a189

Please sign in to comment.