rlgym_learn_backend/common/
misc.rs

1use std::mem::align_of;
2
3use pyo3::{
4    intern,
5    sync::GILOnceCell,
6    types::{PyAnyMethods, PyDict},
7    Bound, IntoPyObject, PyAny, PyErr, PyObject, PyResult, Python,
8};
9use which;
10
11pub fn py_hash(v: &Bound<'_, PyAny>) -> PyResult<i64> {
12    v.call_method0(intern!(v.py(), "__hash__"))?
13        .extract::<i64>()
14}
15
16static INTERNED_CAT: GILOnceCell<PyObject> = GILOnceCell::new();
17static INTERNED_EMPTY: GILOnceCell<PyObject> = GILOnceCell::new();
18
19pub fn get_bytes_to_alignment<T>(addr: usize) -> usize {
20    let alignment = align_of::<T>();
21    let aligned_addr = addr.wrapping_add(alignment - 1) & 0usize.wrapping_sub(alignment);
22    aligned_addr.wrapping_sub(addr)
23}
24
25#[allow(dead_code)]
26pub fn initialize_python() -> pyo3::PyResult<()> {
27    // Due to https://github.com/ContinuumIO/anaconda-issues/issues/11439,
28    // we first need to set PYTHONHOME. To do so, we will look for whatever
29    // directory on PATH currently has python.exe.
30    let python_exe = which::which("python").unwrap();
31    let python_home = python_exe.parent().unwrap();
32    // The Python C API uses null-terminated UTF-16 strings, so we need to
33    // encode the path into that format here.
34    // We could use the Windows FFI modules provided in the standard library,
35    // but we want this to work cross-platform, so we do things more manually.
36    let mut python_home = python_home
37        .to_str()
38        .unwrap()
39        .encode_utf16()
40        .collect::<Vec<u16>>();
41    python_home.push(0);
42    unsafe {
43        pyo3::ffi::Py_SetPythonHome(python_home.as_ptr());
44    }
45    // Once we've set the configuration we need, we can go on and manually
46    // initialize PyO3.
47    pyo3::prepare_freethreaded_python();
48    // Now add cwd to python path
49    Python::with_gil::<_, PyResult<_>>(|py| {
50        Ok(py
51            .import("sys")?
52            .getattr("path")?
53            .call_method1("insert", (0, std::env::current_dir()?.to_str().unwrap()))?
54            .unbind())
55    })?;
56    Ok(())
57}
58
59pub fn clone_list<'py>(py: Python<'py>, list: &Vec<PyObject>) -> Vec<PyObject> {
60    list.iter().map(|obj| obj.clone_ref(py)).collect()
61}
62
63pub fn tensor_slice_1d<'py>(
64    py: Python<'py>,
65    tensor: &Bound<'py, PyAny>,
66    start: usize,
67    stop: usize,
68) -> PyResult<Bound<'py, PyAny>> {
69    Ok(tensor.call_method1(intern!(py, "narrow"), (0, start, stop - start))?)
70}
71
72pub fn torch_cat<'py>(py: Python<'py>, obj: &[&PyObject]) -> PyResult<Bound<'py, PyAny>> {
73    Ok(INTERNED_CAT
74        .get_or_try_init::<_, PyErr>(py, || Ok(py.import("torch")?.getattr("cat")?.unbind()))?
75        .bind(py)
76        .call1((obj,))?)
77}
78
79pub fn torch_empty<'py>(
80    shape: &Bound<'py, PyAny>,
81    dtype: &Bound<'py, PyAny>,
82) -> PyResult<Bound<'py, PyAny>> {
83    let py = shape.py();
84    Ok(INTERNED_EMPTY
85        .get_or_try_init::<_, PyErr>(py, || Ok(py.import("torch")?.getattr("empty")?.unbind()))?
86        .bind(py)
87        .call(
88            (shape,),
89            Some(&PyDict::from_sequence(
90                &vec![(intern!(py, "dtype"), dtype)].into_pyobject(py)?,
91            )?),
92        )?)
93}
rlgym_learn_backend/common/misc.rs

rlgym_learn_backend/common/
misc.rs
