feat: wrap some of starlark::values::Heap

README.md

@@ -50,6 +50,27 @@ version 0.12.x.
     - [x] `AstModule`
 * [ ] `starlark::typing`
 * [ ] `starlark::values`
+    - [ ] `AggregateHeapProfileInfo`
+    - [ ] `Demand`
+    - [ ] `Freezer`
+    - [ ] `FrozenHeap`
+    - [ ] `FrozenHeapRef`
+    - [ ] `FrozenRef`
+    - [ ] `FrozenValue`
+    - [ ] `FrozenValueTyped`
+    - [x] `Heap` -- partially done
+    - [ ] `OwnedFrozenRef`
+    - [ ] `OwnedFrozenValue`
+    - [ ] `OwnedFrozenValueTyped`
+    - [ ] `StarlarkIterator`
+    - [ ] `StarlarkStrNRepr`
+    - [ ] `Tracer`
+    - [ ] `Value`
+    - [ ] `ValueIdentity`
+    - [ ] `ValueOf`
+    - [ ] `ValueOfUnchecked`
+    - [ ] `ValueTyped`
+    - [ ] `ValueTypedComplex`
 
 ## License
 

src/lib.rs

@@ -4,6 +4,7 @@ mod codemap;
 mod environment;
 mod repr_utils;
 mod syntax;
+mod values;
 
 /// A Python module implemented in Rust.
 #[pymodule]
@@ -20,5 +21,7 @@ fn starlark_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<syntax::PyAstModule>()?;
     m.add_class::<syntax::PyDialect>()?;
     m.add_class::<syntax::PyDialectTypes>()?;
+    m.add_class::<values::PyHeap>()?;
+    m.add_class::<values::PyHeapSummary>()?;
     Ok(())
 }

src/values.rs

@@ -0,0 +1,77 @@
+use std::collections::HashMap;
+
+use pyo3::{prelude::*, types::PyTuple};
+use starlark::values::Heap;
+
+/// Information about the data stored on a heap.
+#[pyclass(module = "starlark_pyo3", name = "HeapSummary")]
+pub(crate) struct PyHeapSummary(HashMap<String, (usize, usize)>);
+
+impl From<HashMap<String, (usize, usize)>> for PyHeapSummary {
+    fn from(value: HashMap<String, (usize, usize)>) -> Self {
+        Self(value)
+    }
+}
+
+#[pymethods]
+impl PyHeapSummary {
+    /// (Count, total size) by type.
+    fn summary(slf: PyRef<'_, Self>) -> HashMap<String, Bound<'_, PyTuple>> {
+        let mut x = HashMap::new();
+        for (k, v) in &slf.0 {
+            let v = vec![v.0, v.1];
+            x.insert(k.clone(), PyTuple::new_bound(slf.py(), v));
+        }
+        x
+    }
+
+    /// Total number of bytes allocated.
+    #[getter]
+    fn total_allocated_bytes(&self) -> usize {
+        self.0.values().map(|(_count, bytes)| bytes).sum()
+    }
+}
+
+/// A heap on which `Value`s can be allocated.
+#[pyclass(module = "starlark_pyo3", name = "Heap")]
+pub(crate) struct PyHeap(Heap);
+
+impl From<Heap> for PyHeap {
+    fn from(value: Heap) -> Self {
+        Self(value)
+    }
+}
+
+#[pymethods]
+impl PyHeap {
+    /// Create a new `Heap`.
+    #[new]
+    fn py_new() -> Self {
+        Heap::new().into()
+    }
+
+    /// Number of bytes allocated on this heap, not including any memory
+    /// allocated outside of the starlark heap.
+    #[getter]
+    fn allocated_bytes(&self) -> usize {
+        self.0.allocated_bytes()
+    }
+
+    /// Peak memory allocated to this heap, even if the value is now lower
+    /// as a result of a subsequent garbage collection.
+    #[getter]
+    fn peak_allocated_bytes(&self) -> usize {
+        self.0.peak_allocated_bytes()
+    }
+
+    /// Number of bytes allocated by the heap but not yet filled.
+    #[getter]
+    fn available_bytes(&self) -> usize {
+        self.0.available_bytes()
+    }
+
+    /// Obtain a summary of how much memory is currently allocated by this heap.
+    fn allocated_summary(&self) -> PyHeapSummary {
+        self.0.allocated_summary().summary().into()
+    }
+}

starlark_pyo3.pyi

@@ -152,3 +152,20 @@ class AstModule:
     @property
     def stmt_locations(self) -> list[FileSpan]: ...
     def replace_binary_operators(self, replace: dict[str, str]) -> None: ...
+
+# starlark::values
+
+class HeapSummary:
+    def summary(self) -> dict[str, tuple[int, int]]: ...
+    @property
+    def total_allocated_bytes(self) -> int: ...
+
+class Heap:
+    def __init__(self) -> None: ...
+    @property
+    def allocated_bytes(self) -> int: ...
+    @property
+    def peak_allocated_bytes(self) -> int: ...
+    @property
+    def available_bytes(self) -> int: ...
+    def allocated_summary(self) -> HeapSummary: ...

tests/test_values.py

@@ -0,0 +1,12 @@
+import starlark_pyo3
+
+
+def test_empty_heap():
+    h = starlark_pyo3.Heap()
+    assert h.allocated_bytes == 0
+    assert h.peak_allocated_bytes == 0
+    assert h.available_bytes == 0
+
+    s = h.allocated_summary()
+    assert s.summary() == {}
+    assert s.total_allocated_bytes == 0