polars_python/functions/
io.rs

1use std::io::BufReader;
2
3#[cfg(any(feature = "ipc", feature = "parquet"))]
4use polars::prelude::ArrowSchema;
5use polars::prelude::PlPathRef;
6use polars_io::cloud::CloudOptions;
7use pyo3::prelude::*;
8use pyo3::types::PyDict;
9
10use crate::conversion::Wrap;
11use crate::error::PyPolarsErr;
12use crate::file::{EitherRustPythonFile, get_either_file};
13
14#[cfg(feature = "ipc")]
15#[pyfunction]
16pub fn read_ipc_schema(py: Python<'_>, py_f: PyObject) -> PyResult<Bound<'_, PyDict>> {
17    use arrow::io::ipc::read::read_file_metadata;
18    let metadata = match get_either_file(py_f, false)? {
19        EitherRustPythonFile::Rust(r) => {
20            read_file_metadata(&mut BufReader::new(r)).map_err(PyPolarsErr::from)?
21        },
22        EitherRustPythonFile::Py(mut r) => read_file_metadata(&mut r).map_err(PyPolarsErr::from)?,
23    };
24
25    let dict = PyDict::new(py);
26    fields_to_pydict(&metadata.schema, &dict)?;
27    Ok(dict)
28}
29
30#[cfg(feature = "parquet")]
31#[pyfunction]
32pub fn read_parquet_metadata(
33    py: Python,
34    py_f: PyObject,
35    storage_options: Option<Vec<(String, String)>>,
36    credential_provider: Option<PyObject>,
37    retries: usize,
38) -> PyResult<Bound<PyDict>> {
39    use std::io::Cursor;
40
41    use polars_error::feature_gated;
42    use polars_io::pl_async::get_runtime;
43    use polars_parquet::read::read_metadata;
44    use polars_parquet::read::schema::read_custom_key_value_metadata;
45    use polars_utils::plpath::PlPath;
46
47    use crate::file::{PythonScanSourceInput, get_python_scan_source_input};
48
49    let metadata = match get_python_scan_source_input(py_f, false)? {
50        PythonScanSourceInput::Buffer(buf) => {
51            read_metadata(&mut Cursor::new(buf)).map_err(PyPolarsErr::from)?
52        },
53        PythonScanSourceInput::Path(p) => {
54            let cloud_options = parse_cloud_options(
55                Some(p.as_ref()),
56                storage_options,
57                credential_provider,
58                retries,
59            )?;
60            match p {
61                PlPath::Local(local) => {
62                    let file = polars_utils::open_file(&local).map_err(PyPolarsErr::from)?;
63                    read_metadata(&mut BufReader::new(file)).map_err(PyPolarsErr::from)?
64                },
65                PlPath::Cloud(cloud) => {
66                    use polars::prelude::ParquetObjectStore;
67                    use polars_error::PolarsResult;
68
69                    feature_gated!("cloud", {
70                        get_runtime().block_on(async {
71                            let mut reader = ParquetObjectStore::from_uri(
72                                &cloud.to_string(),
73                                cloud_options.as_ref(),
74                                None,
75                            )
76                            .await?;
77                            let result = reader.get_metadata().await?;
78                            PolarsResult::Ok((**result).clone())
79                        })
80                    })
81                    .map_err(PyPolarsErr::from)?
82                },
83            }
84        },
85        PythonScanSourceInput::File(f) => {
86            read_metadata(&mut BufReader::new(f)).map_err(PyPolarsErr::from)?
87        },
88    };
89
90    let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
91    let dict = PyDict::new(py);
92    for (key, value) in key_value_metadata.into_iter() {
93        dict.set_item(key.as_str(), value.as_str())?;
94    }
95    Ok(dict)
96}
97
98#[cfg(any(feature = "ipc", feature = "parquet"))]
99fn fields_to_pydict(schema: &ArrowSchema, dict: &Bound<'_, PyDict>) -> PyResult<()> {
100    for field in schema.iter_values() {
101        let dt = Wrap(polars::prelude::DataType::from_arrow_field(field));
102        dict.set_item(field.name.as_str(), &dt)?;
103    }
104    Ok(())
105}
106
107#[cfg(feature = "clipboard")]
108#[pyfunction]
109pub fn read_clipboard_string() -> PyResult<String> {
110    use arboard;
111    let mut clipboard =
112        arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
113    let result = clipboard
114        .get_text()
115        .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
116    Ok(result)
117}
118
119#[cfg(feature = "clipboard")]
120#[pyfunction]
121pub fn write_clipboard_string(s: &str) -> PyResult<()> {
122    use arboard;
123    let mut clipboard =
124        arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
125    clipboard
126        .set_text(s)
127        .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
128    Ok(())
129}
130
131pub fn parse_cloud_options<'a>(
132    first_path: Option<PlPathRef<'a>>,
133    storage_options: Option<Vec<(String, String)>>,
134    credential_provider: Option<PyObject>,
135    retries: usize,
136) -> PyResult<Option<CloudOptions>> {
137    let result = if let Some(first_path) = first_path {
138        #[cfg(feature = "cloud")]
139        {
140            use polars_io::cloud::credential_provider::PlCredentialProvider;
141
142            use crate::prelude::parse_cloud_options;
143
144            let first_path_url = first_path.to_str();
145            let cloud_options =
146                parse_cloud_options(first_path_url, storage_options.unwrap_or_default())?;
147
148            Some(
149                cloud_options
150                    .with_max_retries(retries)
151                    .with_credential_provider(
152                        credential_provider.map(PlCredentialProvider::from_python_builder),
153                    ),
154            )
155        }
156
157        #[cfg(not(feature = "cloud"))]
158        {
159            None
160        }
161    } else {
162        None
163    };
164    Ok(result)
165}
polars_python/functions/io.rs

polars_python/functions/
io.rs
