From 9743322aa2620f3541886641201ee5e39dfd3290 Mon Sep 17 00:00:00 2001
From: Noa <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 27 Sep 2021 16:16:09 -0500
Subject: [PATCH] Allow running packages passed to command line

---
 src/lib.rs              | 74 ++++++++++++++++++++++++++++-------------
 vm/src/builtins/dict.rs |  9 +++++
 2 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index ea9a9b74b..2591353f2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -45,13 +45,14 @@ extern crate log;
 
 use clap::{App, AppSettings, Arg, ArgMatches};
 use rustpython_vm::{
-    builtins::PyInt, compile, exceptions::print_exception, match_class, scope::Scope, stdlib::sys,
-    InitParameter, Interpreter, ItemProtocol, PyResult, PySettings, TypeProtocol, VirtualMachine,
+    builtins::PyDictRef, builtins::PyInt, compile, exceptions::print_exception, match_class,
+    scope::Scope, stdlib::sys, InitParameter, Interpreter, ItemProtocol, PyObjectRef, PyResult,
+    PySettings, TryFromObject, TypeProtocol, VirtualMachine,
 };
 
 use std::convert::TryInto;
 use std::env;
-use std::path::PathBuf;
+use std::path::Path;
 use std::process;
 use std::str::FromStr;
 
@@ -613,35 +614,60 @@ fn run_module(vm: &VirtualMachine, module: &str) -> PyResult<()> {
     Ok(())
 }
 
-fn run_script(vm: &VirtualMachine, scope: Scope, script_file: &str) -> PyResult<()> {
-    debug!("Running file {}", script_file);
-    let mut file_path = PathBuf::from(script_file);
-    let file_meta = file_path.metadata().unwrap_or_else(|e| {
-        error!("can't open file '{}': {}", file_path.display(), e);
-        process::exit(1);
-    });
-    if file_meta.is_dir() {
-        file_path.push("__main__.py");
-        if !file_path.is_file() {
-            error!("can't find '__main__' module in '{}'", file_path.display());
-            process::exit(1);
+fn get_importer(path: &str, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
+    let path_importer_cache = vm.get_attribute(vm.sys_module.clone(), "path_importer_cache")?;
+    let path_importer_cache = PyDictRef::try_from_object(vm, path_importer_cache)?;
+    if let Some(importer) = path_importer_cache.get_item_option(path, vm)? {
+        return Ok(Some(importer));
+    }
+    let path = vm.ctx.new_utf8_str(path);
+    let path_hooks = vm.get_attribute(vm.sys_module.clone(), "path_hooks")?;
+    let mut importer = None;
+    for path_hook in vm.extract_elements(&path_hooks)? {
+        match vm.invoke(&path_hook, (path.clone(),)) {
+            Ok(imp) => {
+                importer = Some(imp);
+                break;
+            }
+            Err(e) if e.isinstance(&vm.ctx.exceptions.import_error) => continue,
+            Err(e) => return Err(e),
         }
     }
+    Ok(if let Some(imp) = importer {
+        let imp = path_importer_cache.get_or_insert(vm, path, || imp.clone())?;
+        Some(imp)
+    } else {
+        None
+    })
+}
 
-    let dir = file_path.parent().unwrap().to_str().unwrap().to_owned();
+fn insert_sys_path(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<()> {
     let sys_path = vm.get_attribute(vm.sys_module.clone(), "path").unwrap();
-    vm.call_method(&sys_path, "insert", (0, dir))?;
+    vm.call_method(&sys_path, "insert", (0, obj))?;
+    Ok(())
+}
 
-    match std::fs::read_to_string(&file_path) {
+fn run_script(vm: &VirtualMachine, scope: Scope, script_file: &str) -> PyResult<()> {
+    debug!("Running file {}", script_file);
+    if get_importer(script_file, vm)?.is_some() {
+        insert_sys_path(vm, vm.ctx.new_utf8_str(script_file))?;
+        let runpy = vm.import("runpy", None, 0)?;
+        let run_module_as_main = vm.get_attribute(runpy, "_run_module_as_main")?;
+        vm.invoke(
+            &run_module_as_main,
+            (vm.ctx.new_utf8_str("__main__"), false),
+        )?;
+        return Ok(());
+    }
+    let dir = Path::new(script_file).parent().unwrap().to_str().unwrap();
+    insert_sys_path(vm, vm.ctx.new_utf8_str(dir))?;
+
+    match std::fs::read_to_string(script_file) {
         Ok(source) => {
-            _run_string(vm, scope, &source, file_path.to_str().unwrap().to_owned())?;
+            _run_string(vm, scope, &source, script_file.to_owned())?;
         }
         Err(err) => {
-            error!(
-                "Failed reading file '{}': {:?}",
-                file_path.to_str().unwrap(),
-                err.kind()
-            );
+            error!("Failed reading file '{}': {:?}", script_file, err.kind());
             process::exit(1);
         }
     }
diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs
index e9b5b549b..21e01390c 100644
--- a/vm/src/builtins/dict.rs
+++ b/vm/src/builtins/dict.rs
@@ -297,6 +297,15 @@ impl PyDict {
             .setdefault(vm, key, || default.unwrap_or_none(vm))
     }
 
+    pub fn get_or_insert(
+        &self,
+        vm: &VirtualMachine,
+        key: PyObjectRef,
+        default: impl FnOnce() -> PyObjectRef,
+    ) -> PyResult {
+        self.entries.setdefault(vm, key, default)
+    }
+
     #[pymethod]
     pub fn copy(&self) -> PyDict {
         PyDict {