Skip to content

Commit

Permalink
fix: manually set sys.prefix, exec_prefix and sys.path
Browse files Browse the repository at this point in the history 
When activating a venv in the shell, sys.base_prefix and
sys.base_exec_prefix should be set to the installation location while
sys.prefix and sys.exec_prefix should be set to the venv dir.
Unfortunately, when initialize_venv() and init_pyo3() are called, we
can't use Py_InitializeFromConfig() to set any of these and certain
platforms are unable to find python-build-standalone. For now, we'll
temporarily set PYTHONHOME in init_pyo3() to the installation location
to make python work. By setting it at this point in the code, sys.prefix
and sys.exec_prefix end up also being set to the installation location,
which is fine when not under a venv, but is different from when entering
an venv.

To address this, in PYTHON_INIT.call_once() and when VIRTUAL_ENV is set
(which initialize_venv() will have set at this point), manually set
sys.prefix and sys.exec_prefix to what is in VIRTUAL_ENV.

Similarly, when activating a venv in the shell, sys.path is appended to
have the venv's site-packages dir. Previously we were setting PYTHONPATH
in initialize_venv() which ensures that the venv's site-packages dir is
in sys.path, but this ends up having the venv's site-packages dir first
in sys.path. To correct this, don't set PYTHONPATH any more and instead
adjust PYTHON_INIT.call_once() to append the venv's site-packages dir to
sys.path when VIRTUAL_ENV is set.

Finally, when exiting init_pyo3(), unconditionally unset PYTHONHOME when
VIRTUAL_ENV is set (like activation scripts do) and restore/unset when
it isn't.

Prior to these changes, all target incorrectly had the venv's
site-packages first in sys.path and OSX and Windows additionally had an
incorrect sys.prefix and sys.exec_prefix. With these initialization
changes in place, the runtime environment for the plugins is much closer
to that of a shell activated venv.
  • Loading branch information
@jdstrand
jdstrand committed Feb 25, 2025
1 parent 6ce074b commit 1f8001d
Showing 1 changed file with 86 additions and 34 deletions.
120 changes: 86 additions & 34 deletions influxdb3_processing_engine/src/virtualenv.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,28 +105,19 @@ fn get_python_version() -> Result<(u8, u8), std::io::Error> {
Ok((major, minor))
}

fn set_pythonpath(venv_dir: &Path) -> Result<(), std::io::Error> {
let (major, minor) = get_python_version()?;
let site_packages = if cfg!(target_os = "windows") {
venv_dir.join("Lib").join("site-packages")
} else {
venv_dir
.join("lib")
.join(format!("python{}.{}", major, minor))
.join("site-packages")
};

debug!("Setting PYTHONPATH to: {}", site_packages.to_string_lossy());
unsafe {
std::env::set_var("PYTHONPATH", &site_packages);
}

Ok(())
}

pub fn init_pyo3() {
// PYO3 compiled against python-build-standalone needs PYTHONHOME set
// to the installation location for initialization to work correctly
// When activating a venv in the shell, sys.base_prefix and
// sys.base_exec_prefix should be set to the installation location
// while sys.prefix and sys.exec_prefix should be set to the venv dir.
// Unfortunately, at this point, we can't use Py_InitializeFromConfig()
// to set any of these and certain platforms are unable to find
// python-build-standalone. For now, we'll set PYTHONHOME to the
// installation location to make python work, but by setting it now
// (ie, after Py_InitializeFromConfig() can be called), sys.prefix and
// sys.exec_prefix also end up being set to the installation location.
//
// Note: init_pyo3() is called after initialize_venv() and initialize_venv()
// will set VIRTUAL_ENV as appropriate.
let python_inst = find_python_install();
let orig_pyhome_env = env::var("PYTHONHOME").ok();
if let Some(path) = python_inst {
Expand All @@ -143,28 +134,91 @@ pub fn init_pyo3() {
PYTHON_INIT.call_once(|| {
pyo3::prepare_freethreaded_python();

// This sets the signal handler fo SIGINT to be the default, allowing CTRL+C to work.
// See https://github.com/PyO3/pyo3/issues/3218.
Python::with_gil(|py| {
// This sets the signal handler fo SIGINT to be the default, allowing CTRL+C to work.
// See https://github.com/PyO3/pyo3/issues/3218.
py.run(
&CString::new("import signal;signal.signal(signal.SIGINT, signal.SIG_DFL)")
.unwrap(),
None,
None,
)
.expect("should be able to set signal handler.");

if let Ok(v) = env::var("VIRTUAL_ENV") {
// XXX: ideally we would reorganize the code so that we can
// call Py_InitializeFromConfig() with base_prefix and
// base_exec_prefix set to the python-build-standalone location
// (/path/to/python-build-standalone) and prefix and
// exec_prefix to the value of VIRTUAL_ENV. For now, clean up
// the side-effect of having set PYTHONHOME, above, and set
// sys.prefix and sys.exec_prefix manually.
debug!(
"Updating sys.prefix and sys.exec_prefix to VIRTUAL_ENV: {}",
v
);
let prefix = Path::new(&v);
// Note: need to use python raw strings (r'') for Windows paths
py.run(
&CString::new(format!(
"import sys; sys.prefix = r'{}'; sys.exec_prefix = r'{}'",
prefix.display(),
prefix.display()
))
.unwrap(),
None,
None,
)
.expect("should be able set sys.prefix and sys.exec_prefix");

// When using activate scripts, initializing python within a
// venv results in the venv's site-packages being appended to
// the end of sys.path, so do that here. We could instead set
// the PYTHONPATH env var, but that prepends to sys.path.
if let Ok((major, minor)) = get_python_version() {
let venv_dir = Path::new(&v);
let site_packages = if cfg!(target_os = "windows") {
venv_dir.join("Lib").join("site-packages")
} else {
venv_dir
.join("lib")
.join(format!("python{}.{}", major, minor))
.join("site-packages")
};

debug!("Updating sys.path to append: {}", site_packages.display());
// Note: need to use python raw strings (r'') for Windows paths
py.run(
&CString::new(format!(
"import sys; sys.path.append(r'{}') if r'{}' not in sys.path else None",
site_packages.display(),
site_packages.display(),
))
.unwrap(),
None,
None,
)
.expect("should be able append to sys.path");
}
}
});
});

// now that we're initialized, restore/unset PYTHONHOME
match orig_pyhome_env {
Some(v) => {
debug!("Restoring previous PYTHONHOME to: {}", v);
unsafe { env::set_var("PYTHONHOME", v) }
}
None => {
debug!("Unsetting temporary PYTHONHOME");
unsafe { env::remove_var("PYTHONHOME") }
// Now that we're initialized, unset/restore PYTHONHOME
if env::var("VIRTUAL_ENV").is_ok() {
// venv activate scripts always unset PYTHONHOME, let's do the same
debug!("Unsetting temporary PYTHONHOME when VIRTUAL_ENV is set");
unsafe { env::remove_var("PYTHONHOME") }
} else {
match orig_pyhome_env {
Some(v) => {
debug!("Restoring previous PYTHONHOME to: {}", v);
unsafe { env::set_var("PYTHONHOME", v) }
}
None => {
debug!("Unsetting temporary PYTHONHOME");
unsafe { env::remove_var("PYTHONHOME") }
}
}
}
}
Expand Down Expand Up @@ -223,7 +277,5 @@ pub(crate) fn initialize_venv(venv_path: &Path) -> Result<(), VenvError> {
debug!("VIRTUAL_ENV set to: {}", v);
}

set_pythonpath(venv_path)?;

Ok(())
}

0 comments on commit 1f8001d

Please sign in to comment.