Fix #43

Author: awolverp

.gitignore

@@ -5,3 +5,4 @@ __pycache__
 /.coverage
 /.pytest_cache
 /htmlcov
+/backup

python/tests/mixin.py

@@ -138,7 +138,7 @@ def test_maxmemory_enforced_base_types(self):
         if self.NO_POLICY:
             with pytest.raises(OverflowError):
                 cache[10] = 10
-            
+
             assert 1 in cache
         else:
             cache[10] = 10
@@ -560,3 +560,11 @@ def test_copy(self):
 
         c2.insert(1, 1)
         assert 1 not in c1
+
+    def test_cache_type(self):
+        class AType:
+            pass
+
+        cache = self.CACHE(maxsize=0, **self.KWARGS)
+        cache[AType] = AType
+        assert cache[AType] is AType

src/common.rs

@@ -1,3 +1,5 @@
+use pyo3::types::PyAnyMethods;
+
 macro_rules! non_zero_or {
     ($num:expr, $_else:expr) => {
         unsafe { core::num::NonZeroUsize::new_unchecked(if $num == 0 { $_else } else { $num }) }
@@ -180,23 +182,36 @@ macro_rules! extract_pickle_tuple {
 #[inline]
 #[cfg(not(PyPy))]
 pub fn pyobject_size(py: pyo3::Python<'_>, obj: &pyo3::Py<pyo3::PyAny>) -> pyo3::PyResult<usize> {
-    use pyo3::types::PyAnyMethods;
+    static SIZEOF_METHOD_NAME: &'static std::ffi::CStr = c"__sizeof__";
+
+    // PyPy does not support __sizeof__ or sys.getsizeof
+    let sizeof_method = obj.bind(py).getattr(SIZEOF_METHOD_NAME)?;
 
-    obj.bind(py).call_method0("__sizeof__")?.extract::<usize>()
+    unsafe {
+        if pyo3::ffi::PyType_Check(obj.as_ptr()) == 1 {
+            sizeof_method.call1((obj,))?.extract::<usize>()
+        } else {
+            sizeof_method.call0()?.extract::<usize>()
+        }
+    }
 }
 
 #[inline]
 #[cfg(PyPy)]
 pub fn pyobject_size(py: pyo3::Python<'_>, obj: &pyo3::Py<pyo3::PyAny>) -> pyo3::PyResult<usize> {
-    use pyo3::types::PyAnyMethods;
-
-    static SIZEOF_METHOD_NAME: &'static str = "__sizeof__";
+    static SIZEOF_METHOD_NAME: &'static std::ffi::CStr = c"__sizeof__";
 
     // PyPy does not support __sizeof__ or sys.getsizeof
-    let has_sizeof = obj.bind(py).getattr_opt(SIZEOF_METHOD_NAME)?;
+    let sizeof_method = obj.bind(py).getattr_opt(SIZEOF_METHOD_NAME)?;
 
-    match has_sizeof {
-        Some(x) => x.call0()?.extract::<usize>(),
+    match sizeof_method {
+        Some(sizeof_method) => unsafe {
+            if pyo3::ffi::PyType_Check(obj.as_ptr()) == 1 {
+                sizeof_method.call1((obj,))?.extract::<usize>()
+            } else {
+                sizeof_method.call0()?.extract::<usize>()
+            }
+        },
         None => Ok(1),
     }
 }