polars_python/functions/
io.rs1use std::io::BufReader;
2
3use arrow::array::Utf8ViewArray;
4#[cfg(any(feature = "ipc", feature = "parquet"))]
5use polars::prelude::ArrowSchema;
6use polars_core::datatypes::create_enum_dtype;
7use pyo3::prelude::*;
8use pyo3::types::PyDict;
9
10use crate::conversion::Wrap;
11use crate::error::PyPolarsErr;
12use crate::file::{EitherRustPythonFile, get_either_file};
13use crate::prelude::ArrowDataType;
14
15#[cfg(feature = "ipc")]
16#[pyfunction]
17pub fn read_ipc_schema(py: Python<'_>, py_f: PyObject) -> PyResult<Bound<PyDict>> {
18 use arrow::io::ipc::read::read_file_metadata;
19 let metadata = match get_either_file(py_f, false)? {
20 EitherRustPythonFile::Rust(r) => {
21 read_file_metadata(&mut BufReader::new(r)).map_err(PyPolarsErr::from)?
22 },
23 EitherRustPythonFile::Py(mut r) => read_file_metadata(&mut r).map_err(PyPolarsErr::from)?,
24 };
25
26 let dict = PyDict::new(py);
27 fields_to_pydict(&metadata.schema, &dict)?;
28 Ok(dict)
29}
30
31#[cfg(feature = "parquet")]
32#[pyfunction]
33pub fn read_parquet_schema(py: Python<'_>, py_f: PyObject) -> PyResult<Bound<PyDict>> {
34 use polars_parquet::read::{infer_schema, read_metadata};
35
36 let metadata = match get_either_file(py_f, false)? {
37 EitherRustPythonFile::Rust(r) => {
38 read_metadata(&mut BufReader::new(r)).map_err(PyPolarsErr::from)?
39 },
40 EitherRustPythonFile::Py(mut r) => read_metadata(&mut r).map_err(PyPolarsErr::from)?,
41 };
42 let arrow_schema = infer_schema(&metadata).map_err(PyPolarsErr::from)?;
43
44 let dict = PyDict::new(py);
45 fields_to_pydict(&arrow_schema, &dict)?;
46 Ok(dict)
47}
48
49#[cfg(feature = "parquet")]
50#[pyfunction]
51pub fn read_parquet_metadata(py: Python, py_f: PyObject) -> PyResult<Bound<PyDict>> {
52 use polars_parquet::read::read_metadata;
53 use polars_parquet::read::schema::read_custom_key_value_metadata;
54
55 let metadata = match get_either_file(py_f, false)? {
56 EitherRustPythonFile::Rust(r) => {
57 read_metadata(&mut BufReader::new(r)).map_err(PyPolarsErr::from)?
58 },
59 EitherRustPythonFile::Py(mut r) => read_metadata(&mut r).map_err(PyPolarsErr::from)?,
60 };
61
62 let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
63 let dict = PyDict::new(py);
64 for (key, value) in key_value_metadata.into_iter() {
65 dict.set_item(key.as_str(), value.as_str())?;
66 }
67 Ok(dict)
68}
69
70#[cfg(any(feature = "ipc", feature = "parquet"))]
71fn fields_to_pydict(schema: &ArrowSchema, dict: &Bound<'_, PyDict>) -> PyResult<()> {
72 for field in schema.iter_values() {
73 let dt = if field.is_enum() {
74 Wrap(create_enum_dtype(Utf8ViewArray::new_empty(
75 ArrowDataType::Utf8View,
76 )))
77 } else {
78 Wrap(polars::prelude::DataType::from_arrow_field(field))
79 };
80 dict.set_item(field.name.as_str(), &dt)?;
81 }
82 Ok(())
83}
84
85#[cfg(feature = "clipboard")]
86#[pyfunction]
87pub fn read_clipboard_string() -> PyResult<String> {
88 use arboard;
89 let mut clipboard =
90 arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
91 let result = clipboard
92 .get_text()
93 .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
94 Ok(result)
95}
96
97#[cfg(feature = "clipboard")]
98#[pyfunction]
99pub fn write_clipboard_string(s: &str) -> PyResult<()> {
100 use arboard;
101 let mut clipboard =
102 arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
103 clipboard
104 .set_text(s)
105 .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
106 Ok(())
107}