lib.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.

// Suppress clippy::redundant_closure warning from pyo3 generated code
#![allow(clippy::redundant_closure)]

use std::collections::HashMap;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::os::raw::c_int;
use std::str::FromStr;

use ::opendal as od;
use pyo3::create_exception;
use pyo3::exceptions::PyException;
use pyo3::exceptions::PyFileExistsError;
use pyo3::exceptions::PyFileNotFoundError;
use pyo3::exceptions::PyIOError;
use pyo3::exceptions::PyNotImplementedError;
use pyo3::exceptions::PyPermissionError;
use pyo3::exceptions::PyValueError;
use pyo3::ffi;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::AsPyPointer;

mod asyncio;
mod layers;

use crate::asyncio::*;

create_exception!(opendal, Error, PyException, "OpenDAL related errors");

/// A bytes-like object that implements buffer protocol.
#[pyclass(module = "opendal")]
struct Buffer {
    inner: Vec<u8>,
}

#[pymethods]
impl Buffer {
    unsafe fn __getbuffer__(
        slf: PyRefMut<Self>,
        view: *mut ffi::Py_buffer,
        flags: c_int,
    ) -> PyResult<()> {
        let bytes = slf.inner.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(())
    }
}

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

fn add_layers(mut op: od::Operator, layers: Vec<layers::Layer>) -> PyResult<od::Operator> {
    for layer in layers {
        match layer {
            layers::Layer::Retry(layers::RetryLayer(inner)) => op = op.layer(inner),
            layers::Layer::ImmutableIndex(layers::ImmutableIndexLayer(inner)) => {
                op = op.layer(inner)
            }
            layers::Layer::ConcurrentLimit(layers::ConcurrentLimitLayer(inner)) => {
                op = op.layer(inner)
            }
        }
    }
    Ok(op)
}

fn build_operator(
    scheme: od::Scheme,
    map: HashMap<String, String>,
    layers: Vec<layers::Layer>,
    blocking: bool,
) -> PyResult<od::Operator> {
    let mut op = od::Operator::via_map(scheme, map).map_err(format_pyerr)?;
    if blocking && !op.info().full_capability().blocking {
        let runtime = pyo3_asyncio::tokio::get_runtime();
        let _guard = runtime.enter();
        op = op.layer(od::layers::BlockingLayer::create().expect("blocking layer must be created"));
    }

    add_layers(op, layers)
}

/// `Operator` is the entry for all public blocking APIs
///
/// Create a new blocking `Operator` with the given `scheme` and options(`**kwargs`).
#[pyclass(module = "opendal")]
struct Operator(od::BlockingOperator);

#[pymethods]
impl Operator {
    #[new]
    #[pyo3(signature = (scheme, *, layers=Vec::new(), **map))]
    pub fn new(scheme: &str, layers: Vec<layers::Layer>, map: Option<&PyDict>) -> PyResult<Self> {
        let scheme = od::Scheme::from_str(scheme)
            .map_err(|err| {
                od::Error::new(od::ErrorKind::Unexpected, "unsupported scheme").set_source(err)
            })
            .map_err(format_pyerr)?;
        let map = map
            .map(|v| {
                v.extract::<HashMap<String, String>>()
                    .expect("must be valid hashmap")
            })
            .unwrap_or_default();

        Ok(Operator(
            build_operator(scheme, map, layers, true)?.blocking(),
        ))
    }

    /// Read the whole path into bytes.
    pub fn read<'p>(&'p self, py: Python<'p>, path: &str) -> PyResult<&'p PyAny> {
        let buffer = self
            .0
            .read(path)
            .map_err(format_pyerr)
            .map(Buffer::from)?
            .into_py(py);
        let memoryview =
            unsafe { py.from_owned_ptr_or_err(ffi::PyMemoryView_FromObject(buffer.as_ptr()))? };
        Ok(memoryview)
    }

    /// Open a file-like reader for the given path.
    pub fn open_reader(&self, path: &str) -> PyResult<Reader> {
        self.0
            .reader(path)
            .map(|reader| Reader(Some(reader)))
            .map_err(format_pyerr)
    }

    /// Write bytes into given path.
    #[pyo3(signature = (path, bs, **kwargs))]
    pub fn write(&self, path: &str, bs: Vec<u8>, kwargs: Option<&PyDict>) -> PyResult<()> {
        let opwrite = build_opwrite(kwargs)?;
        let mut write = self.0.write_with(path, bs).append(opwrite.append());
        if let Some(buffer) = opwrite.buffer() {
            write = write.buffer(buffer);
        }
        if let Some(content_type) = opwrite.content_type() {
            write = write.content_type(content_type);
        }
        if let Some(content_disposition) = opwrite.content_disposition() {
            write = write.content_disposition(content_disposition);
        }
        if let Some(cache_control) = opwrite.cache_control() {
            write = write.cache_control(cache_control);
        }

        write.call().map_err(format_pyerr)
    }

    /// Get current path's metadata **without cache** directly.
    pub fn stat(&self, path: &str) -> PyResult<Metadata> {
        self.0.stat(path).map_err(format_pyerr).map(Metadata)
    }

    /// Create a dir at given path.
    ///
    /// # Notes
    ///
    /// To indicate that a path is a directory, it is compulsory to include
    /// a trailing / in the path. Failure to do so may result in
    /// `NotADirectory` error being returned by OpenDAL.
    ///
    /// # Behavior
    ///
    /// - Create on existing dir will succeed.
    /// - Create dir is always recursive, works like `mkdir -p`
    pub fn create_dir(&self, path: &str) -> PyResult<()> {
        self.0.create_dir(path).map_err(format_pyerr)
    }

    /// Delete given path.
    ///
    /// # Notes
    ///
    /// - Delete not existing error won't return errors.
    pub fn delete(&self, path: &str) -> PyResult<()> {
        self.0.delete(path).map_err(format_pyerr)
    }

    /// List current dir path.
    pub fn list(&self, path: &str) -> PyResult<BlockingLister> {
        Ok(BlockingLister(self.0.lister(path).map_err(format_pyerr)?))
    }

    /// List dir in flat way.
    pub fn scan(&self, path: &str) -> PyResult<BlockingLister> {
        Ok(BlockingLister(
            self.0
                .lister_with(path)
                .delimiter("")
                .call()
                .map_err(format_pyerr)?,
        ))
    }

    fn __repr__(&self) -> String {
        let info = self.0.info();
        let name = info.name();
        if name.is_empty() {
            format!("Operator(\"{}\", root=\"{}\")", info.scheme(), info.root())
        } else {
            format!(
                "Operator(\"{}\", root=\"{}\", name=\"{name}\")",
                info.scheme(),
                info.root()
            )
        }
    }
}

/// A file-like blocking reader.
/// Can be used as a context manager.
#[pyclass(module = "opendal")]
struct Reader(Option<od::BlockingReader>);

impl Reader {
    fn as_mut(&mut self) -> PyResult<&mut od::BlockingReader> {
        let reader = self
            .0
            .as_mut()
            .ok_or_else(|| PyValueError::new_err("I/O operation on closed file."))?;
        Ok(reader)
    }
}

#[pymethods]
impl Reader {
    /// Read and return size bytes, or if size is not given, until EOF.
    #[pyo3(signature = (size=None,))]
    pub fn read<'p>(&'p mut self, py: Python<'p>, size: Option<usize>) -> PyResult<&'p PyAny> {
        let reader = self.as_mut()?;
        let buffer = match size {
            Some(size) => {
                let mut buffer = vec![0; size];
                reader
                    .read_exact(&mut buffer)
                    .map_err(|err| PyIOError::new_err(err.to_string()))?;
                buffer
            }
            None => {
                let mut buffer = Vec::new();
                reader
                    .read_to_end(&mut buffer)
                    .map_err(|err| PyIOError::new_err(err.to_string()))?;
                buffer
            }
        };
        let buffer = Buffer::from(buffer).into_py(py);
        let memoryview =
            unsafe { py.from_owned_ptr_or_err(ffi::PyMemoryView_FromObject(buffer.as_ptr()))? };
        Ok(memoryview)
    }

    /// `Reader` doesn't support write.
    /// Raises a `NotImplementedError` if called.
    pub fn write(&mut self, _bs: &[u8]) -> PyResult<()> {
        Err(PyNotImplementedError::new_err(
            "Reader does not support write",
        ))
    }

    /// Change the stream position to the given byte offset.
    /// offset is interpreted relative to the position indicated by `whence`.
    /// The default value for whence is `SEEK_SET`. Values for `whence` are:
    ///
    /// * `SEEK_SET` or `0` – start of the stream (the default); offset should be zero or positive
    /// * `SEEK_CUR` or `1` – current stream position; offset may be negative
    /// * `SEEK_END` or `2` – end of the stream; offset is usually negative
    ///
    /// Return the new absolute position.
    #[pyo3(signature = (pos, whence = 0))]
    pub fn seek(&mut self, pos: i64, whence: u8) -> PyResult<u64> {
        let whence = match whence {
            0 => SeekFrom::Start(pos as u64),
            1 => SeekFrom::Current(pos),
            2 => SeekFrom::End(pos),
            _ => return Err(PyValueError::new_err("invalid whence")),
        };
        let reader = self.as_mut()?;
        reader
            .seek(whence)
            .map_err(|err| PyIOError::new_err(err.to_string()))
    }

    /// Return the current stream position.
    pub fn tell(&mut self) -> PyResult<u64> {
        let reader = self.as_mut()?;
        reader
            .stream_position()
            .map_err(|err| PyIOError::new_err(err.to_string()))
    }

    pub fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    pub fn __exit__(&mut self, _exc_type: PyObject, _exc_value: PyObject, _traceback: PyObject) {
        drop(self.0.take());
    }
}

#[pyclass(unsendable, module = "opendal")]
struct BlockingLister(od::BlockingLister);

#[pymethods]
impl BlockingLister {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }
    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<PyObject>> {
        match slf.0.next() {
            Some(Ok(entry)) => Ok(Some(Entry(entry).into_py(slf.py()))),
            Some(Err(err)) => {
                let pyerr = format_pyerr(err);
                Err(pyerr)
            }
            None => Ok(None),
        }
    }
}

#[pyclass(module = "opendal")]
struct Entry(od::Entry);

#[pymethods]
impl Entry {
    /// Path of entry. Path is relative to operator's root.
    #[getter]
    pub fn path(&self) -> &str {
        self.0.path()
    }

    fn __str__(&self) -> &str {
        self.0.path()
    }

    fn __repr__(&self) -> String {
        format!("Entry({:?})", self.0.path())
    }
}

#[pyclass(module = "opendal")]
struct Metadata(od::Metadata);

#[pymethods]
impl Metadata {
    #[getter]
    pub fn content_disposition(&self) -> Option<&str> {
        self.0.content_disposition()
    }

    /// Content length of this entry.
    #[getter]
    pub fn content_length(&self) -> u64 {
        self.0.content_length()
    }

    /// Content MD5 of this entry.
    #[getter]
    pub fn content_md5(&self) -> Option<&str> {
        self.0.content_md5()
    }

    /// Content Type of this entry.
    #[getter]
    pub fn content_type(&self) -> Option<&str> {
        self.0.content_type()
    }

    /// ETag of this entry.
    #[getter]
    pub fn etag(&self) -> Option<&str> {
        self.0.etag()
    }

    /// mode represent this entry's mode.
    #[getter]
    pub fn mode(&self) -> EntryMode {
        EntryMode(self.0.mode())
    }
}

#[pyclass(module = "opendal")]
struct EntryMode(od::EntryMode);

#[pymethods]
impl EntryMode {
    /// Returns `True` if this is a file.
    pub fn is_file(&self) -> bool {
        self.0.is_file()
    }

    /// Returns `True` if this is a directory.
    pub fn is_dir(&self) -> bool {
        self.0.is_dir()
    }

    pub fn __repr__(&self) -> &'static str {
        match self.0 {
            od::EntryMode::FILE => "EntryMode.FILE",
            od::EntryMode::DIR => "EntryMode.DIR",
            od::EntryMode::Unknown => "EntryMode.UNKNOWN",
        }
    }
}

#[pyclass(module = "opendal")]
struct PresignedRequest(od::raw::PresignedRequest);

#[pymethods]
impl PresignedRequest {
    /// Return the URL of this request.
    #[getter]
    pub fn url(&self) -> String {
        self.0.uri().to_string()
    }

    /// Return the HTTP method of this request.
    #[getter]
    pub fn method(&self) -> &str {
        self.0.method().as_str()
    }

    /// Return the HTTP headers of this request.
    #[getter]
    pub fn headers(&self) -> PyResult<HashMap<&str, &str>> {
        let mut headers = HashMap::new();
        for (k, v) in self.0.header().iter() {
            let k = k.as_str();
            let v = v.to_str().map_err(|err| Error::new_err(err.to_string()))?;
            if headers.insert(k, v).is_some() {
                return Err(Error::new_err("duplicate header"));
            }
        }
        Ok(headers)
    }
}

fn format_pyerr(err: od::Error) -> PyErr {
    use od::ErrorKind::*;
    match err.kind() {
        NotFound => PyFileNotFoundError::new_err(err.to_string()),
        AlreadyExists => PyFileExistsError::new_err(err.to_string()),
        PermissionDenied => PyPermissionError::new_err(err.to_string()),
        Unsupported => PyNotImplementedError::new_err(err.to_string()),
        _ => Error::new_err(err.to_string()),
    }
}

/// recognize OpWrite-equivalent options passed as python dict
pub(crate) fn build_opwrite(kwargs: Option<&PyDict>) -> PyResult<od::raw::OpWrite> {
    use od::raw::OpWrite;
    let mut op = OpWrite::new();

    let dict = if let Some(kwargs) = kwargs {
        kwargs
    } else {
        return Ok(op);
    };

    if let Some(append) = dict.get_item("append") {
        let v = append
            .extract::<bool>()
            .map_err(|err| PyValueError::new_err(format!("append must be bool, got {}", err)))?;
        op = op.with_append(v);
    }

    if let Some(buffer) = dict.get_item("buffer") {
        let v = buffer
            .extract::<usize>()
            .map_err(|err| PyValueError::new_err(format!("buffer must be usize, got {}", err)))?;
        op = op.with_buffer(v);
    }

    if let Some(content_type) = dict.get_item("content_type") {
        let v = content_type.extract::<String>().map_err(|err| {
            PyValueError::new_err(format!("content_type must be str, got {}", err))
        })?;
        op = op.with_content_type(v.as_str());
    }

    if let Some(content_disposition) = dict.get_item("content_disposition") {
        let v = content_disposition.extract::<String>().map_err(|err| {
            PyValueError::new_err(format!("content_disposition must be str, got {}", err))
        })?;
        op = op.with_content_disposition(v.as_str());
    }

    if let Some(cache_control) = dict.get_item("cache_control") {
        let v = cache_control.extract::<String>().map_err(|err| {
            PyValueError::new_err(format!("cache_control must be str, got {}", err))
        })?;
        op = op.with_cache_control(v.as_str());
    }

    Ok(op)
}

/// OpenDAL Python binding
///
/// ## Installation
///
/// ```bash
/// pip install opendal
/// ```
///
/// ## Usage
///
/// ```python
/// import opendal
///
/// op = opendal.Operator("fs", root="/tmp")
/// op.write("test.txt", b"Hello World")
/// print(op.read("test.txt"))
/// print(op.stat("test.txt").content_length)
/// ```
///
/// Or using the async API:
///
/// ```python
/// import asyncio
///
/// async def main():
/// op = opendal.AsyncOperator("fs", root="/tmp")
/// await op.write("test.txt", b"Hello World")
/// print(await op.read("test.txt"))
///
/// asyncio.run(main())
/// ```
#[pymodule]
fn _opendal(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Operator>()?;
    m.add_class::<Reader>()?;
    m.add_class::<AsyncOperator>()?;
    m.add_class::<AsyncReader>()?;
    m.add_class::<Entry>()?;
    m.add_class::<EntryMode>()?;
    m.add_class::<Metadata>()?;
    m.add_class::<PresignedRequest>()?;
    m.add("Error", py.get_type::<Error>())?;

    let layers = layers::create_submodule(py)?;
    m.add_submodule(layers)?;
    py.import("sys")?
        .getattr("modules")?
        .set_item("opendal.layers", layers)?;

    Ok(())
}