1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

// Code is copied and modified from the original source code: https://github.com/apache/incubator-opendal/blob/main/bindings/python/src/utils.rs
//
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::os::raw::c_int;

use pyo3::ffi;
use pyo3::prelude::*;

pub(crate) const INLINED_U8: usize = 64;

pub type Bytes = smallvec::SmallVec<[u8; INLINED_U8]>;

/// A bytes-like object that implements buffer protocol.
#[pyclass]
pub struct Buffer(Bytes);

impl From<Bytes> for Buffer {
  fn from(inner: Bytes) -> Self {
    Buffer(inner)
  }
}

impl From<Vec<u8>> for Buffer {
  fn from(inner: Vec<u8>) -> Self {
    Buffer(inner.into())
  }
}

impl From<&[u8]> for Buffer {
  fn from(inner: &[u8]) -> Self {
    Buffer(inner.into())
  }
}

impl From<&str> for Buffer {
  fn from(inner: &str) -> Self {
    Buffer(inner.as_bytes().into())
  }
}

impl From<String> for Buffer {
  fn from(inner: String) -> Self {
    Buffer(inner.into_bytes().into())
  }
}

impl From<Buffer> for Bytes {
  fn from(buffer: Buffer) -> Self {
    buffer.0
  }
}

impl Buffer {
  pub fn new(inner: Bytes) -> Self {
    Buffer(inner)
  }

  /// Consume self to build a memory view
  pub fn into_memory_view(self, py: Python) -> PyResult<Py<PyAny>> {
    let buffer = self.into_py(py);

    unsafe { PyObject::from_owned_ptr_or_err(py, ffi::PyMemoryView_FromObject(buffer.as_ptr())) }
  }

  /// Consume self to build a memory view ref.
  pub fn into_memory_view_ref(self, py: Python) -> PyResult<&PyAny> {
    let buffer = self.into_py(py);
    let view = unsafe { py.from_owned_ptr_or_err(ffi::PyMemoryView_FromObject(buffer.as_ptr()))? };

    Ok(view)
  }
}

#[pymethods]
impl Buffer {
  unsafe fn __getbuffer__(
    slf: PyRefMut<Self>,
    view: *mut ffi::Py_buffer,
    flags: c_int,
  ) -> PyResult<()> {
    let bytes = slf.0.as_slice();
    let ret = ffi::PyBuffer_FillInfo(
      view,
      slf.as_ptr() as *mut _,
      bytes.as_ptr() as *mut _,
      bytes.len().try_into().unwrap(),
      1, // read only
      flags,
    );
    if ret == -1 {
      return Err(PyErr::fetch(slf.py()));
    }
    Ok(())
  }
}