Implement PYTHONPATH

OddBloke · OddBloke · commit da34b86a9c81 · 2018-08-24T22:55:30.000-04:00
This modifies the import logic to use sys.path, and uses PYTHONPATH to
generate the initial contents of sys.path.
diff --git a/tests/snippets/import.py b/tests/snippets/import.py
@@ -13,6 +13,12 @@
 assert import_target.X == aliased_func()
 assert import_target.Y == aliased_other_func()
 
+# TODO: Once we can determine current directory, use that to construct this
+# path:
+#import sys
+#sys.path.append("snippets/import_directory")
+#import nested_target
+
 #try:
 #    X
 #except NameError:
diff --git a/tests/test_snippets.py b/tests/test_snippets.py
@@ -78,12 +78,9 @@ def run_via_rustpython(filename, test_type):
     log_level = 'info' if test_type == _TestType.benchmark else 'trace'
     env['RUST_LOG'] = '{},cargo=error,jobserver=error'.format(log_level)
     env['RUST_BACKTRACE'] = '1'
-    # XXX: Once we support PYTHONPATH (or similar), we should use that instead
-    # of changing directory
-    cwd = os.path.dirname(filename)
-    with pushd(RUSTPYTHON_RUNNER_DIR):
-        subprocess.check_call(
-            ['cargo', 'run', '--release', filename], env=env, cwd=cwd)
+    env['PYTHONPATH'] = os.path.dirname(filename)
+    subprocess.check_call(
+        ['cargo', 'run', '--release', filename], env=env)
 
 
 def create_test_function(cls, filename, method, test_type):
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -21,7 +21,8 @@ fn import_module(vm: &mut VirtualMachine, module: &String) -> PyResult {
     }
 
     // Time to search for module in any place:
-    let filepath = find_source(module).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
+    let filepath =
+        find_source(vm, module).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
     let source = parser::read_file(filepath.as_path())
         .map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
 
@@ -62,14 +63,29 @@ pub fn import(vm: &mut VirtualMachine, module: &String, symbol: &Option<String>)
     Ok(obj)
 }
 
-fn find_source(name: &String) -> io::Result<PathBuf> {
+fn find_source(vm: &VirtualMachine, name: &String) -> io::Result<PathBuf> {
+    let sys_path = vm.sys_module.get_item("path").unwrap();
+    let paths: Vec<PathBuf> = match sys_path.borrow().kind {
+        PyObjectKind::List { ref elements } => elements
+            .iter()
+            .filter_map(|item| match item.borrow().kind {
+                PyObjectKind::String { ref value } => Some(PathBuf::from(value)),
+                _ => None,
+            }).collect(),
+        _ => panic!("sys.path unexpectedly not a list"),
+    };
+
     let suffixes = [".py", "/__init__.py"];
-    let filepaths = suffixes
-        .iter()
-        .map(|suffix| format!("{}{}", name, suffix))
-        .map(|filename| PathBuf::from(filename));
+    let mut filepaths = vec![];
+    for path in paths {
+        for suffix in suffixes.iter() {
+            let mut filepath = path.clone();
+            filepath.push(format!("{}{}", name, suffix));
+            filepaths.push(filepath);
+        }
+    }
 
-    match filepaths.filter(|p| p.exists()).next() {
+    match filepaths.iter().filter(|p| p.exists()).next() {
         Some(path) => Ok(path.to_path_buf()),
         None => Err(io::Error::new(
             NotFound,
diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -3,11 +3,19 @@
 /*
  * The magic sys module.
  */
+use std::env;
 
 use super::pyobject::{DictProtocol, PyContext, PyObjectRef};
 
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
-    let path = ctx.new_list(vec![]);
+    let mut path_list = match env::var_os("PYTHONPATH") {
+        Some(paths) => env::split_paths(&paths)
+            .map(|path| ctx.new_str(path.to_str().unwrap().to_string()))
+            .collect(),
+        None => vec![],
+    };
+    path_list.insert(0, ctx.new_str("".to_string()));
+    let path = ctx.new_list(path_list);
     let modules = ctx.new_dict();
     let sys_name = "sys".to_string();
     let sys_mod = ctx.new_module(&sys_name, ctx.new_scope(None));
