Mercurial > public > mercurial-scm > hg

annotate rust/hg-pyo3/src/util.rs @ 52822:4f41a8acf350

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
rust-pyo3: add a `with_pybytes_buffer` util This is very similar to the one we have in `hg-cpython`, and serves the same purpose. Explanations inline.
author Rapha?l Gom?s <rgomes@octobus.net>
date Fri, 03 Jan 2025 12:43:52 +0100
parents 42b219a1404a
children dd3a2948804f
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
52822
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
1 use hg::errors::HgError;
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
2 use hg::revlog::inner_revlog::RevisionBuffer;
52780
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
3 use pyo3::buffer::{Element, PyBuffer};
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
4 use pyo3::exceptions::PyValueError;
52407
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
5 use pyo3::prelude::*;
52822
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
6 use pyo3::types::{PyBytes, PyDict};
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
7
52407
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
8 /// Create the module, with `__package__` given from parent
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
9 ///
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
10 /// According to PyO3 documentation, which links to
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
11 /// <https://github.com/PyO3/pyo3/issues/1517>, the same convoluted
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
12 /// write to sys.modules has to be made as with the `cpython` crate.
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
13 pub(crate) fn new_submodule<'py>(
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
14 py: Python<'py>,
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
15 package_name: &str,
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
16 name: &str,
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
17 ) -> PyResult<Bound<'py, PyModule>> {
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
18 let dotted_name = &format!("{}.{}", package_name, name);
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
19 let m = PyModule::new(py, name)?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
20 m.add("__package__", package_name)?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
21 m.add("__doc__", "DAG operations - Rust implementation")?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
22
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
23 let sys = PyModule::import(py, "sys")?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
24 // according to the doc, we could make a static PyString out of
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
25 // "modules" with the `intern!` macro, but this is used only at
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
26 // registration so it may not be worth the effort.
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
27 let sys_modules: Bound<'_, PyDict> = sys.getattr("modules")?.extract()?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
28 sys_modules.set_item(dotted_name, &m)?;
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
29 // Example C code (see pyexpat.c and import.c) will "give away the
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
30 // reference", but we won't because it will be consumed once the
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
31 // Rust PyObject is dropped.
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
32 Ok(m)
c5128c541021 rust-pyo3: facility for submodule registration, using it for dagop
Georges Racinet <georges.racinet@cloudcrane.io>
parents:
diff changeset
33 }
52780
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
34
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
35 /// Type shortcut for the kind of bytes slice trait objects that are used in
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
36 /// particular for mmap data
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
37 type BoxedBytesSlice =
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
38 Box<dyn std::ops::Deref<Target = [u8]> + Send + Sync + 'static>;
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
39
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
40 /// Take a Python object backed by a Python buffer, and return the underlying
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
41 /// [`PyBuffer`] along with the Rust slice into said buffer.
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
42 ///
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
43 /// The caller needs to make sure that the Python buffer is not freed before
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
44 /// the slice, otherwise we'd get a dangling pointer once the incoming
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
45 /// object is freed from Python side. This can be achieved by storing it a
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
46 /// Python object.
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
47 ///
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
48 /// The typical use case is to extract mmap data to make it useable in the
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
49 /// constructs from the `hg` crate.
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
50 ///
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
51 /// # Safety
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
52 ///
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
53 /// The caller must make sure that the incoming Python object is kept around
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
54 /// for at least as long as the returned [`BoxedBytesSlice`].
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
55 // TODO in PyO3, we already get a reference with two lifetimes, and we
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
56 // could even take a `Borrowed<'a, 'py, T>`.
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
57 // So perhaps we could tie everything together with a lifetime so that is
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
58 // is, after all, safe, and this could be called something like `share_buffer`.
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
59 #[deny(unsafe_op_in_unsafe_fn)]
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
60 pub unsafe fn take_buffer_with_slice(
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
61 data: &Bound<'_, PyAny>,
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
62 ) -> PyResult<(PyBuffer<u8>, BoxedBytesSlice)> {
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
63 let buf = PyBuffer::<u8>::get(data)?;
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
64 let len = buf.item_count();
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
65
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
66 // Build a slice from the buffer data
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
67 let cbuf = buf.buf_ptr();
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
68 let bytes = if std::mem::size_of::<u8>() == buf.item_size()
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
69 && buf.is_c_contiguous()
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
70 && u8::is_compatible_format(buf.format())
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
71 && buf.dimensions() == 1
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
72 && buf.readonly()
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
73 {
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
74 unsafe { std::slice::from_raw_parts(cbuf as *const u8, len) }
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
75 } else {
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
76 return Err(PyValueError::new_err(
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
77 "buffer has an invalid memory representation",
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
78 ));
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
79 };
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
80
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
81 Ok((buf, Box::new(bytes)))
42b219a1404a rust-pyo3-revlog: InnerRevlog definition and constructor
Georges Racinet <georges.racinet@cloudcrane.io>
parents: 52407
diff changeset
82 }
52822
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
83
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
84 /// Takes an initialization function `init` which writes bytes to a
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
85 /// Python-backed buffer, to save on a (potentially large) memory allocation
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
86 /// and copy. If `init` fails to write the full expected length `len`, an error
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
87 /// is raised.
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
88 pub fn with_pybytes_buffer<F>(
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
89 py: Python,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
90 len: usize,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
91 init: F,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
92 ) -> Result<Py<PyBytes>, hg::revlog::RevlogError>
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
93 where
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
94 F: FnOnce(
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
95 &mut dyn RevisionBuffer<Target = Py<PyBytes>>,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
96 ) -> Result<(), hg::revlog::RevlogError>,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
97 {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
98 // Largely inspired by code in PyO3
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
99 // https://pyo3.rs/main/doc/pyo3/types/struct.pybytes#method.new_bound_with
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
100 unsafe {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
101 let pyptr = pyo3::ffi::PyBytes_FromStringAndSize(
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
102 std::ptr::null(),
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
103 len as pyo3::ffi::Py_ssize_t,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
104 );
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
105 let pybytes = Bound::from_owned_ptr_or_err(py, pyptr)
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
106 .map_err(|e| HgError::abort_simple(e.to_string()))?
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
107 .downcast_into_unchecked();
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
108 let buffer: *mut u8 = pyo3::ffi::PyBytes_AsString(pyptr).cast();
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
109 debug_assert!(!buffer.is_null());
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
110 let mut rev_buf = PyRevisionBuffer::new(pybytes.unbind(), buffer, len);
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
111 // Initialise the bytestring in init
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
112 // If init returns an Err, the buffer is deallocated by `pybytes`
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
113 init(&mut rev_buf).map(|_| rev_buf.finish())
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
114 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
115 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
116
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
117 /// Wrapper around a Python-provided buffer into which the revision contents
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
118 /// will be written. Done for speed in order to save a large allocation + copy.
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
119 struct PyRevisionBuffer {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
120 py_bytes: Py<PyBytes>,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
121 _buf: *mut u8,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
122 len: usize,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
123 current_buf: *mut u8,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
124 current_len: usize,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
125 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
126
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
127 impl PyRevisionBuffer {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
128 /// # Safety
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
129 ///
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
130 /// `buf` should be the start of the allocated bytes of `bytes`, and `len`
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
131 /// exactly the length of said allocated bytes.
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
132 #[inline]
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
133 unsafe fn new(bytes: Py<PyBytes>, buf: *mut u8, len: usize) -> Self {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
134 Self {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
135 py_bytes: bytes,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
136 _buf: buf,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
137 len,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
138 current_len: 0,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
139 current_buf: buf,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
140 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
141 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
142
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
143 /// Number of bytes that have been copied to. Will be different to the
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
144 /// total allocated length of the buffer unless the revision is done being
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
145 /// written.
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
146 #[inline]
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
147 fn current_len(&self) -> usize {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
148 self.current_len
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
149 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
150 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
151
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
152 impl RevisionBuffer for PyRevisionBuffer {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
153 type Target = Py<PyBytes>;
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
154
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
155 #[inline]
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
156 fn extend_from_slice(&mut self, slice: &[u8]) {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
157 assert!(self.current_len + slice.len() <= self.len);
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
158 unsafe {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
159 // We cannot use `copy_from_nonoverlapping` since it's *possible*
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
160 // to create a slice from the same Python memory region using
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
161 // [`PyBytesDeref`]. Probable that LLVM has an optimization anyway?
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
162 self.current_buf.copy_from(slice.as_ptr(), slice.len());
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
163 self.current_buf = self.current_buf.add(slice.len());
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
164 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
165 self.current_len += slice.len()
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
166 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
167
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
168 #[inline]
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
169 fn finish(self) -> Self::Target {
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
170 // catch unzeroed bytes before it becomes undefined behavior
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
171 assert_eq!(
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
172 self.current_len(),
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
173 self.len,
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
174 "not enough bytes read for revision"
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
175 );
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
176 self.py_bytes
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
177 }
4f41a8acf350 rust-pyo3: add a `with_pybytes_buffer` util
Rapha?l Gom?s <rgomes@octobus.net>
parents: 52780
diff changeset
178 }