Created
November 23, 2021 21:01
-
-
Save adriangb/e90af7603fc89da94cb8754f32835325 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use std::cmp; | |
use std::collections::HashMap; | |
use std::collections::hash_map::Entry; | |
use std::hash; | |
use pyo3::prelude::*; | |
use pyo3::types::{PyTuple}; | |
// We can't put a Py<PyAny> directly into a HashMap key | |
// So to be able to hold references to arbitrary Python objects in HashMap as keys | |
// we wrap them in a struct that gets the hash() when it receives the object from Python | |
// and then just echoes back that hash when called Rust needs to hash it | |
#[derive(Debug)] | |
struct HashedAny { | |
o: Py<PyAny>, | |
hash: isize, | |
} | |
impl <'source>FromPyObject<'source> for HashedAny | |
{ | |
fn extract(ob: &'source PyAny) -> PyResult<Self> { | |
Ok(HashedAny{ o: ob.into(), hash: ob.hash()? }) | |
} | |
} | |
impl hash::Hash for HashedAny { | |
fn hash<H: hash::Hasher>(&self, state: &mut H) { | |
state.write(&self.hash.to_be_bytes()) | |
} | |
} | |
impl cmp::PartialEq for HashedAny { | |
fn eq(&self, other: &Self) -> bool { | |
self.o.eq(&other.o) | |
} | |
} | |
impl cmp::Eq for HashedAny {} | |
#[pyfunction(args="*")] | |
fn func(ob: HashedAny, args: &PyTuple) -> PyResult<()> { | |
let mut map: HashMap<HashedAny, usize> = HashMap::new(); | |
map.insert(ob, 0); | |
let mut idx: usize = 1; | |
for item in args.iter() { | |
match map.entry(HashedAny::extract(&item)?) { | |
Entry::Occupied(_) => (), | |
Entry::Vacant(e) => { | |
idx += 1; | |
e.insert(idx); | |
} | |
} | |
} | |
println!("{:?}", map); | |
Ok(()) | |
} | |
#[pymodule] | |
fn di_lib(_py: Python, m: &PyModule) -> PyResult<()> { | |
m.add_function(wrap_pyfunction!(func, m)?)?; | |
Ok(()) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment