#[cfg(feature = "python")]
use pyo3::exceptions::{PyRuntimeError, PyValueError};
#[cfg(feature = "python")]
use pyo3::prelude::*;
#[cfg(feature = "python")]
use pyo3::types::PyModule;
#[cfg(feature = "python")]
use std::mem;
#[cfg(feature = "python")]
use std::sync::Mutex;
#[cfg(feature = "python")]
use crate::{JSONTools, JsonOutput};
#[cfg(feature = "python")]
use crate::json_parser;
#[cfg(feature = "python")]
use pythonize::{depythonize, pythonize};
#[cfg(feature = "python")]
pyo3::create_exception!(
json_tools_rs,
JsonToolsError,
pyo3::exceptions::PyException,
"Python exception for JSON Tools operations"
);
#[cfg(feature = "python")]
#[inline]
fn lock_config(mutex: &Mutex<JSONTools>) -> PyResult<std::sync::MutexGuard<'_, JSONTools>> {
mutex
.lock()
.map_err(|e| PyRuntimeError::new_err(format!("internal config lock poisoned: {e}")))
}
#[cfg(feature = "python")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DataFrameType {
Pandas,
Polars,
PyArrow,
PySpark,
Generic,
}
#[cfg(feature = "python")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SeriesType {
Pandas,
Polars,
PyArrow,
PySpark,
Generic,
}
#[cfg(feature = "python")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DataStructureType {
DataFrame(DataFrameType),
Series(SeriesType),
}
#[cfg(feature = "python")]
#[pyclass(name = "JsonOutput", frozen, from_py_object)]
#[derive(Debug, Clone)]
pub struct PyJsonOutput {
inner: JsonOutput,
}
#[cfg(feature = "python")]
#[pymethods]
impl PyJsonOutput {
#[getter]
fn is_single(&self) -> bool {
matches!(self.inner, JsonOutput::Single(_))
}
#[getter]
fn is_multiple(&self) -> bool {
matches!(self.inner, JsonOutput::Multiple(_))
}
fn get_single(&self) -> PyResult<String> {
match &self.inner {
JsonOutput::Single(result) => Ok(result.clone()),
JsonOutput::Multiple(_) => Err(PyValueError::new_err(
"Result contains multiple JSON strings, use get_multiple() instead",
)),
}
}
fn get_multiple(&self) -> PyResult<Vec<String>> {
match &self.inner {
JsonOutput::Single(_) => Err(PyValueError::new_err(
"Result contains single JSON string, use get_single() instead",
)),
JsonOutput::Multiple(results) => Ok(results.clone()),
}
}
fn to_python(&self, py: Python) -> PyResult<Py<PyAny>> {
match &self.inner {
JsonOutput::Single(result) => Ok(result.into_pyobject(py)?.into_any().unbind()),
JsonOutput::Multiple(results) => Ok(results.into_pyobject(py)?.into_any().unbind()),
}
}
fn __repr__(&self) -> String {
match &self.inner {
JsonOutput::Single(result) => format!("JsonOutput.Single('{}')", result),
JsonOutput::Multiple(results) => format!("JsonOutput.Multiple({:?})", results),
}
}
fn __str__(&self) -> String {
match &self.inner {
JsonOutput::Single(result) => result.clone(),
JsonOutput::Multiple(results) => format!("{:?}", results),
}
}
}
#[cfg(feature = "python")]
impl From<JsonOutput> for PyJsonOutput {
fn from(output: JsonOutput) -> Self {
PyJsonOutput { inner: output }
}
}
#[cfg(feature = "python")]
impl PyJsonOutput {
pub fn from_rust_output(output: JsonOutput) -> Self {
PyJsonOutput { inner: output }
}
}
#[cfg(feature = "python")]
#[pyclass(name = "JSONTools")]
pub struct PyJSONTools {
inner: Mutex<JSONTools>,
}
#[cfg(feature = "python")]
impl Default for PyJSONTools {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "python")]
fn detect_data_structure(obj: &Bound<'_, PyAny>) -> PyResult<Option<DataStructureType>> {
if let Some(df_type) = detect_dataframe_type(obj)? {
return Ok(Some(DataStructureType::DataFrame(df_type)));
}
if let Some(series_type) = detect_series_type(obj)? {
return Ok(Some(DataStructureType::Series(series_type)));
}
Ok(None)
}
#[cfg(feature = "python")]
fn detect_dataframe_type(obj: &Bound<'_, PyAny>) -> PyResult<Option<DataFrameType>> {
let module = obj
.getattr("__class__")?
.getattr("__module__")?
.extract::<String>()
.unwrap_or_default();
let class_name = obj
.getattr("__class__")?
.getattr("__name__")?
.extract::<String>()
.unwrap_or_default();
if module.starts_with("pyarrow") && class_name == "Table" {
return Ok(Some(DataFrameType::PyArrow));
}
let has_to_dict = obj.hasattr("to_dict")?;
let has_columns = obj.hasattr("columns")?;
if !has_to_dict && !has_columns {
return Ok(None); }
if class_name != "DataFrame" {
return Ok(None);
}
if module.starts_with("pandas") {
Ok(Some(DataFrameType::Pandas))
} else if module.starts_with("polars") {
Ok(Some(DataFrameType::Polars))
} else if module.contains("pyspark.sql") {
Ok(Some(DataFrameType::PySpark))
} else if has_to_dict || obj.hasattr("to_json")? {
Ok(Some(DataFrameType::Generic))
} else {
Ok(None)
}
}
#[cfg(feature = "python")]
fn detect_series_type(obj: &Bound<'_, PyAny>) -> PyResult<Option<SeriesType>> {
let module = obj
.getattr("__class__")?
.getattr("__module__")?
.extract::<String>()
.unwrap_or_default();
let class_name = obj
.getattr("__class__")?
.getattr("__name__")?
.extract::<String>()
.unwrap_or_default();
if module.starts_with("pyarrow")
&& (class_name.contains("Array") || obj.hasattr("to_pylist")?)
{
return Ok(Some(SeriesType::PyArrow));
}
let has_to_list = obj.hasattr("to_list")? || obj.hasattr("tolist")?;
let has_dtype = obj.hasattr("dtype")?;
if !has_to_list && !has_dtype {
return Ok(None); }
if class_name != "Series" {
return Ok(None);
}
if module.starts_with("pandas") {
Ok(Some(SeriesType::Pandas))
} else if module.starts_with("polars") {
Ok(Some(SeriesType::Polars))
} else if module.contains("pyspark") {
Ok(Some(SeriesType::PySpark))
} else if has_to_list {
Ok(Some(SeriesType::Generic))
} else {
Ok(None)
}
}
#[cfg(feature = "python")]
fn dataframe_to_records(
df: &Bound<'_, PyAny>,
df_type: DataFrameType,
) -> PyResult<Vec<serde_json::Value>> {
let py = df.py();
match df_type {
DataFrameType::Pandas => {
let kwargs = pyo3::types::PyDict::new(py);
kwargs.set_item("orient", "records")?;
let records = df.call_method("to_dict", (), Some(&kwargs))?;
let list = records.cast::<pyo3::types::PyList>()?;
convert_pylist_to_json_values(list)
}
DataFrameType::Polars => {
let records = df.call_method0("to_dicts")?;
let list = records.cast::<pyo3::types::PyList>()?;
convert_pylist_to_json_values(list)
}
DataFrameType::PyArrow => {
let records = df.call_method0("to_pylist")?;
let list = records.cast::<pyo3::types::PyList>()?;
convert_pylist_to_json_values(list)
}
DataFrameType::PySpark => {
let pandas_df = df.call_method0("toPandas")?;
let kwargs = pyo3::types::PyDict::new(py);
kwargs.set_item("orient", "records")?;
let records = pandas_df.call_method("to_dict", (), Some(&kwargs))?;
let list = records.cast::<pyo3::types::PyList>()?;
convert_pylist_to_json_values(list)
}
DataFrameType::Generic => {
if df.hasattr("to_dict")? {
let kwargs = pyo3::types::PyDict::new(py);
kwargs.set_item("orient", "records")?;
let records = df.call_method("to_dict", (), Some(&kwargs))?;
let list = records.cast::<pyo3::types::PyList>()?;
convert_pylist_to_json_values(list)
} else {
Err(JsonToolsError::new_err(
"Generic DataFrame must have to_dict() method",
))
}
}
}
}
#[cfg(feature = "python")]
fn convert_pylist_to_json_values(
list: &Bound<'_, pyo3::types::PyList>,
) -> PyResult<Vec<serde_json::Value>> {
let mut values = Vec::with_capacity(list.len());
for item in list.iter() {
let value: serde_json::Value = depythonize(&item)
.map_err(|e| JsonToolsError::new_err(format!("Failed to convert record: {}", e)))?;
values.push(value);
}
Ok(values)
}
#[cfg(feature = "python")]
fn series_to_list<'py>(
series: &Bound<'py, PyAny>,
series_type: SeriesType,
) -> PyResult<Bound<'py, pyo3::types::PyList>> {
match series_type {
SeriesType::Pandas => {
if series.hasattr("to_list")? {
let list = series.call_method0("to_list")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
} else {
let list = series.call_method0("tolist")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
}
}
SeriesType::Polars => {
let list = series.call_method0("to_list")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
}
SeriesType::PyArrow => {
let list = series.call_method0("to_pylist")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
}
SeriesType::PySpark => {
let pandas_series = series.call_method0("toPandas")?;
let list = pandas_series.call_method0("tolist")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
}
SeriesType::Generic => {
if series.hasattr("to_list")? {
let list = series.call_method0("to_list")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
} else if series.hasattr("tolist")? {
let list = series.call_method0("tolist")?;
Ok(list.cast::<pyo3::types::PyList>()?.clone())
} else {
Err(JsonToolsError::new_err(
"Generic Series must have to_list() or tolist() method",
))
}
}
}
}
#[cfg(feature = "python")]
macro_rules! py_builder_method {
($slf:expr, $tools:ident, $body:expr) => {{
let mut guard = lock_config(&$slf.inner)?;
let $tools = mem::take(&mut *guard);
*guard = $body;
drop(guard);
Ok($slf)
}};
}
#[cfg(feature = "python")]
#[pymethods]
impl PyJSONTools {
#[new]
#[pyo3(text_signature = "()")]
pub fn new() -> Self {
Self {
inner: Mutex::new(JSONTools::new()),
}
}
#[pyo3(text_signature = "($self)")]
#[inline]
pub fn flatten(slf: PyRef<'_, Self>) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.flatten())
}
#[pyo3(text_signature = "($self)")]
#[inline]
pub fn unflatten(slf: PyRef<'_, Self>) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.unflatten())
}
#[pyo3(text_signature = "($self)")]
#[inline]
pub fn normal(slf: PyRef<'_, Self>) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.normal())
}
#[pyo3(text_signature = "($self, separator)")]
#[inline]
pub fn separator(slf: PyRef<'_, Self>, separator: String) -> PyResult<PyRef<'_, Self>> {
if separator.is_empty() {
return Err(PyValueError::new_err("Separator cannot be empty"));
}
py_builder_method!(slf, tools, tools.separator(separator))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn lowercase_keys(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.lowercase_keys(value))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn remove_empty_strings(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.remove_empty_strings(value))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn remove_nulls(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.remove_nulls(value))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn remove_empty_objects(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.remove_empty_objects(value))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn remove_empty_arrays(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.remove_empty_arrays(value))
}
#[pyo3(text_signature = "($self, find, replace)")]
#[inline]
pub fn key_replacement(
slf: PyRef<'_, Self>,
find: String,
replace: String,
) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.key_replacement(find, replace))
}
#[pyo3(text_signature = "($self, find, replace)")]
#[inline]
pub fn value_replacement(
slf: PyRef<'_, Self>,
find: String,
replace: String,
) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.value_replacement(find, replace))
}
#[pyo3(text_signature = "($self, value)")]
#[inline]
pub fn handle_key_collision(slf: PyRef<'_, Self>, value: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.handle_key_collision(value))
}
#[pyo3(text_signature = "($self, enable)")]
#[inline]
pub fn auto_convert_types(slf: PyRef<'_, Self>, enable: bool) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.auto_convert_types(enable))
}
#[pyo3(text_signature = "($self, threshold)")]
#[inline]
pub fn parallel_threshold(slf: PyRef<'_, Self>, threshold: usize) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.parallel_threshold(threshold))
}
#[pyo3(text_signature = "($self, num_threads)")]
#[inline]
pub fn num_threads(
slf: PyRef<'_, Self>,
num_threads: Option<usize>,
) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.num_threads(num_threads))
}
#[pyo3(text_signature = "($self, threshold)")]
#[inline]
pub fn nested_parallel_threshold(
slf: PyRef<'_, Self>,
threshold: usize,
) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.nested_parallel_threshold(threshold))
}
#[pyo3(text_signature = "($self, max)")]
#[inline]
pub fn max_array_index(slf: PyRef<'_, Self>, max: usize) -> PyResult<PyRef<'_, Self>> {
py_builder_method!(slf, tools, tools.max_array_index(max))
}
#[pyo3(text_signature = "($self, json_input)")]
pub fn execute(&self, json_input: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
let py = json_input.py();
match detect_data_structure(json_input)? {
Some(DataStructureType::DataFrame(df_type)) => {
return self.execute_dataframe(json_input, df_type);
}
Some(DataStructureType::Series(series_type)) => {
return self.execute_series(json_input, series_type);
}
None => {
}
}
if let Ok(json_str) = json_input.extract::<String>() {
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_str.as_str());
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process JSON string: {}", e))
})?;
match result {
JsonOutput::Single(processed) => {
Ok(processed.into_pyobject(py)?.into_any().unbind())
}
JsonOutput::Multiple(_) => Err(PyValueError::new_err(
"Unexpected multiple results for single JSON input",
)),
}
} else if json_input.is_instance_of::<pyo3::types::PyDict>() {
let value: serde_json::Value = depythonize(json_input).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert Python dict: {}", e))
})?;
let json_str = json_parser::to_string(&value)
.map_err(|e| JsonToolsError::new_err(format!("Failed to serialize: {}", e)))?;
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_str.as_str());
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process Python dict: {}", e))
})?;
match result {
JsonOutput::Single(processed) => {
let result_value: serde_json::Value = json_parser::from_str(&processed)
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to parse result: {}", e))
})?;
let python_dict = pythonize(py, &result_value).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert to Python: {}", e))
})?;
Ok(python_dict.unbind())
}
JsonOutput::Multiple(_) => Err(PyValueError::new_err(
"Unexpected multiple results for single dict input",
)),
}
} else if json_input.is_instance_of::<pyo3::types::PyList>() {
let list = json_input.cast::<pyo3::types::PyList>()?;
if list.is_empty() {
return Ok(Vec::<String>::new().into_pyobject(py)?.into_any().unbind());
}
let mut json_strings: Vec<String> = Vec::with_capacity(list.len());
let mut is_str_flags: Vec<bool> = Vec::with_capacity(list.len());
let mut has_other_types = false;
for item in list.iter() {
if let Ok(json_str) = item.extract::<String>() {
json_strings.push(json_str);
is_str_flags.push(true);
} else if item.is_instance_of::<pyo3::types::PyDict>() {
let value: serde_json::Value = depythonize(&item).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert dict in list: {}", e))
})?;
let json_str = json_parser::to_string(&value).map_err(|e| {
JsonToolsError::new_err(format!("Failed to serialize dict: {}", e))
})?;
json_strings.push(json_str);
is_str_flags.push(false);
} else {
has_other_types = true;
break;
}
}
if has_other_types {
return Err(PyValueError::new_err(
"List items must be either JSON strings or Python dictionaries",
));
}
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_strings);
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process JSON list: {}", e))
})?;
match result {
JsonOutput::Single(_) => Err(PyValueError::new_err(
"Unexpected single result for multiple input",
)),
JsonOutput::Multiple(processed_list) => {
let all_strings = is_str_flags.iter().all(|&b| b);
let all_dicts = is_str_flags.iter().all(|&b| !b);
if all_strings {
Ok(processed_list.into_pyobject(py)?.into_any().unbind())
} else if all_dicts {
let mut dict_results: Vec<Py<PyAny>> =
Vec::with_capacity(processed_list.len());
for processed_json in processed_list {
let result_value: serde_json::Value =
json_parser::from_str(&processed_json).map_err(|e| {
JsonToolsError::new_err(format!("Failed to parse JSON: {}", e))
})?;
let python_dict = pythonize(py, &result_value).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to convert to Python dict: {}",
e
))
})?;
dict_results.push(python_dict.unbind());
}
Ok(dict_results.into_pyobject(py)?.into_any().unbind())
} else {
let mut mixed_results: Vec<Py<PyAny>> =
Vec::with_capacity(processed_list.len());
for (processed_json, is_str) in processed_list.into_iter().zip(is_str_flags)
{
if is_str {
mixed_results
.push(processed_json.into_pyobject(py)?.into_any().unbind());
} else {
let result_value: serde_json::Value =
json_parser::from_str(&processed_json).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to parse JSON: {}",
e
))
})?;
let python_dict = pythonize(py, &result_value).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to convert to Python dict: {}",
e
))
})?;
mixed_results.push(python_dict.unbind());
}
}
Ok(mixed_results.into_pyobject(py)?.into_any().unbind())
}
}
}
} else {
Err(PyValueError::new_err(
"json_input must be a JSON string, Python dict, list of JSON strings, or list of Python dicts",
))
}
}
#[pyo3(text_signature = "($self, json_input)")]
pub fn execute_to_output(&self, json_input: &Bound<'_, PyAny>) -> PyResult<PyJsonOutput> {
let py = json_input.py();
if let Ok(json_str) = json_input.extract::<String>() {
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_str.as_str());
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process JSON string: {}", e))
})?;
return Ok(PyJsonOutput::from_rust_output(result));
}
if json_input.is_instance_of::<pyo3::types::PyDict>() {
let value: serde_json::Value = depythonize(json_input).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert Python dict: {}", e))
})?;
let json_str = json_parser::to_string(&value)
.map_err(|e| JsonToolsError::new_err(format!("Failed to serialize: {}", e)))?;
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_str.as_str());
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process Python dict: {}", e))
})?;
return Ok(PyJsonOutput::from_rust_output(result));
}
if json_input.is_instance_of::<pyo3::types::PyList>() {
let list = json_input.cast::<pyo3::types::PyList>()?;
if list.is_empty() {
return Ok(PyJsonOutput::from_rust_output(JsonOutput::Multiple(vec![])));
}
let mut json_strings: Vec<String> = Vec::with_capacity(list.len());
for item in list.iter() {
if let Ok(json_str) = item.extract::<String>() {
json_strings.push(json_str);
} else if item.is_instance_of::<pyo3::types::PyDict>() {
let value: serde_json::Value = depythonize(&item).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert dict in list: {}", e))
})?;
let json_str = json_parser::to_string(&value).map_err(|e| {
JsonToolsError::new_err(format!("Failed to serialize dict: {}", e))
})?;
json_strings.push(json_str);
} else {
return Err(PyValueError::new_err(
"List items must be either JSON strings or Python dictionaries",
));
}
}
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_strings);
*guard = tools;
result
})
.map_err(|e| {
JsonToolsError::new_err(format!("Failed to process JSON list: {}", e))
})?;
return Ok(PyJsonOutput::from_rust_output(result));
}
Err(PyValueError::new_err(
"json_input must be a JSON string, Python dict, DataFrame, Series, list of JSON strings, or list of Python dicts",
))
}
}
#[cfg(feature = "python")]
fn reconstruct_dataframe(
py: Python,
df_type: DataFrameType,
processed_dicts: Vec<Py<PyAny>>,
) -> PyResult<Py<PyAny>> {
match df_type {
DataFrameType::Pandas => reconstruct_pandas_df(py, processed_dicts),
DataFrameType::Polars => reconstruct_polars_df(py, processed_dicts),
DataFrameType::PyArrow => reconstruct_pyarrow_table(py, processed_dicts),
DataFrameType::PySpark => {
Ok(processed_dicts.into_pyobject(py)?.into_any().unbind())
}
DataFrameType::Generic => {
Ok(processed_dicts.into_pyobject(py)?.into_any().unbind())
}
}
}
#[cfg(feature = "python")]
fn reconstruct_pandas_df(py: Python, records: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = records.into_pyobject(py)?.into_any().unbind();
match py.import("pandas") {
Ok(pandas) => match pandas.call_method1("DataFrame", (py_list.clone_ref(py),)) {
Ok(df) => Ok(df.unbind()),
Err(_) => Ok(py_list),
},
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
fn reconstruct_polars_df(py: Python, records: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = records.into_pyobject(py)?.into_any().unbind();
match py.import("polars") {
Ok(polars) => match polars.call_method1("DataFrame", (py_list.clone_ref(py),)) {
Ok(df) => Ok(df.unbind()),
Err(_) => Ok(py_list),
},
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
fn reconstruct_pyarrow_table(py: Python, records: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = records.into_pyobject(py)?.into_any().unbind();
match py.import("pyarrow") {
Ok(pyarrow) => {
let table_class = pyarrow.getattr("Table")?;
match table_class.call_method1("from_pylist", (py_list.clone_ref(py),)) {
Ok(table) => Ok(table.unbind()),
Err(_) => Ok(py_list),
}
}
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
fn reconstruct_series(
py: Python,
series_type: SeriesType,
processed_items: Vec<Py<PyAny>>,
) -> PyResult<Py<PyAny>> {
match series_type {
SeriesType::Pandas => reconstruct_pandas_series(py, processed_items),
SeriesType::Polars => reconstruct_polars_series(py, processed_items),
SeriesType::PyArrow => reconstruct_pyarrow_array(py, processed_items),
SeriesType::PySpark => {
Ok(processed_items.into_pyobject(py)?.into_any().unbind())
}
SeriesType::Generic => {
Ok(processed_items.into_pyobject(py)?.into_any().unbind())
}
}
}
#[cfg(feature = "python")]
fn reconstruct_pandas_series(py: Python, items: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = items.into_pyobject(py)?.into_any().unbind();
match py.import("pandas") {
Ok(pandas) => match pandas.call_method1("Series", (py_list.clone_ref(py),)) {
Ok(series) => Ok(series.unbind()),
Err(_) => Ok(py_list),
},
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
fn reconstruct_polars_series(py: Python, items: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = items.into_pyobject(py)?.into_any().unbind();
match py.import("polars") {
Ok(polars) => match polars.call_method1("Series", (py_list.clone_ref(py),)) {
Ok(series) => Ok(series.unbind()),
Err(_) => Ok(py_list),
},
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
fn reconstruct_pyarrow_array(py: Python, items: Vec<Py<PyAny>>) -> PyResult<Py<PyAny>> {
let py_list = items.into_pyobject(py)?.into_any().unbind();
match py.import("pyarrow") {
Ok(pyarrow) => match pyarrow.call_method1("array", (py_list.clone_ref(py),)) {
Ok(array) => Ok(array.unbind()),
Err(_) => Ok(py_list),
},
Err(_) => Ok(py_list),
}
}
#[cfg(feature = "python")]
impl PyJSONTools {
fn execute_dataframe(
&self,
df: &Bound<'_, PyAny>,
df_type: DataFrameType,
) -> PyResult<Py<PyAny>> {
let py = df.py();
let records = dataframe_to_records(df, df_type)?;
let mut json_strings = Vec::with_capacity(records.len());
for record in records {
let json_str = json_parser::to_string(&record).map_err(|e| {
JsonToolsError::new_err(format!("Failed to serialize record: {}", e))
})?;
json_strings.push(json_str);
}
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_strings); *guard = tools;
result
})
.map_err(|e| JsonToolsError::new_err(format!("Failed to process DataFrame: {}", e)))?;
match result {
JsonOutput::Multiple(processed_list) => {
let mut processed_dicts: Vec<Py<PyAny>> = Vec::with_capacity(processed_list.len());
for json_str in processed_list {
let value: serde_json::Value =
json_parser::from_str(&json_str).map_err(|e| {
JsonToolsError::new_err(format!("Failed to parse result: {}", e))
})?;
let py_dict = pythonize(py, &value).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert to Python: {}", e))
})?;
processed_dicts.push(py_dict.unbind());
}
reconstruct_dataframe(py, df_type, processed_dicts)
}
JsonOutput::Single(_) => Err(PyValueError::new_err(
"Unexpected single result for DataFrame input",
)),
}
}
fn execute_series(
&self,
series: &Bound<'_, PyAny>,
series_type: SeriesType,
) -> PyResult<Py<PyAny>> {
let py = series.py();
let list = series_to_list(series, series_type)?;
let mut json_strings: Vec<String> = Vec::with_capacity(list.len());
let mut is_str_flags: Vec<bool> = Vec::with_capacity(list.len());
let mut has_other_types = false;
for item in list.iter() {
if let Ok(json_str) = item.extract::<String>() {
json_strings.push(json_str);
is_str_flags.push(true);
} else if item.is_instance_of::<pyo3::types::PyDict>() {
let value: serde_json::Value = depythonize(&item).map_err(|e| {
JsonToolsError::new_err(format!("Failed to convert dict in series: {}", e))
})?;
let json_str = json_parser::to_string(&value).map_err(|e| {
JsonToolsError::new_err(format!("Failed to serialize dict: {}", e))
})?;
json_strings.push(json_str);
is_str_flags.push(false);
} else {
has_other_types = true;
break;
}
}
if has_other_types {
return Err(PyValueError::new_err(
"Series items must be either JSON strings or Python dictionaries",
));
}
let result = py
.detach(|| {
let mut guard = lock_config(&self.inner)?;
let tools = mem::take(&mut *guard);
let result = tools.execute(json_strings);
*guard = tools;
result
})
.map_err(|e| JsonToolsError::new_err(format!("Failed to process Series: {}", e)))?;
match result {
JsonOutput::Multiple(processed_list) => {
let all_strings = is_str_flags.iter().all(|&b| b);
let all_dicts = is_str_flags.iter().all(|&b| !b);
let processed_items: Vec<Py<PyAny>> = if all_strings {
processed_list
.into_iter()
.map(|s| {
s.into_pyobject(py)
.map(|o| o.into_any().unbind())
.map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to convert string to Python object: {}",
e
))
})
})
.collect::<PyResult<Vec<_>>>()?
} else if all_dicts {
let mut dict_results = Vec::with_capacity(processed_list.len());
for processed_json in processed_list {
let result_value: serde_json::Value =
json_parser::from_str(&processed_json).map_err(|e| {
JsonToolsError::new_err(format!("Failed to parse result: {}", e))
})?;
let python_dict = pythonize(py, &result_value).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to convert to Python dict: {}",
e
))
})?;
dict_results.push(python_dict.unbind());
}
dict_results
} else {
let mut mixed_results = Vec::with_capacity(processed_list.len());
for (processed_json, is_str) in processed_list.into_iter().zip(is_str_flags) {
if is_str {
mixed_results
.push(processed_json.into_pyobject(py)?.into_any().unbind());
} else {
let result_value: serde_json::Value =
json_parser::from_str(&processed_json).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to parse result: {}",
e
))
})?;
let python_dict = pythonize(py, &result_value).map_err(|e| {
JsonToolsError::new_err(format!(
"Failed to convert to Python: {}",
e
))
})?;
mixed_results.push(python_dict.unbind());
}
}
mixed_results
};
reconstruct_series(py, series_type, processed_items)
}
JsonOutput::Single(_) => Err(PyValueError::new_err(
"Unexpected single result for Series input",
)),
}
}
}
#[cfg(feature = "python")]
#[pymodule]
fn json_tools_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyJSONTools>()?;
m.add_class::<PyJsonOutput>()?;
m.add("JsonToolsError", m.py().get_type::<JsonToolsError>())?;
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
m.add("__author__", "JSON Tools RS Contributors")?;
m.add(
"__description__",
"Python bindings for JSON Tools RS - Unified JSON manipulation library with advanced collision handling and filtering",
)?;
Ok(())
}