wingfoil-python 8.0.0

//! Python bindings for the CSV adapter.

use crate::py_element::PyElement;
use crate::py_stream::PyStream;
use anyhow::Context;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::cell::RefCell;
use std::collections::HashMap;
use std::io::Write;
use std::rc::Rc;
use wingfoil::adapters::csv::csv_read;
use wingfoil::{Burst, NanoTime, Node, Stream, StreamOperators};

type CsvRow = HashMap<String, String>;

/// Read a CSV file into a stream of dicts.
///
/// Each tick yields a dict where keys are column headers and values are strings.
/// Rows sharing the same timestamp are collapsed to one-per-tick.
///
/// Args:
///     path: Path to the CSV file (headers required)
///     time_column: Name of the column containing the timestamp
///                  (parsed as integer nanoseconds since epoch)
#[pyfunction]
pub fn py_csv_read(path: String, time_column: String) -> PyResult<PyStream> {
    let stream: Rc<dyn Stream<Burst<CsvRow>>> = csv_read(
        &path,
        move |row: &CsvRow| {
            row.get(&time_column)
                .and_then(|s| s.parse::<u64>().ok())
                .map(NanoTime::new)
                .unwrap_or(NanoTime::ZERO)
        },
        true,
    )
    .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}")))?;

    let py_stream = stream.collapse().map(|row: CsvRow| {
        Python::attach(|py| {
            let dict = PyDict::new(py);
            for (key, value) in &row {
                dict.set_item(key, value)
                    .expect("invariant: inserting str keys/values into a dict cannot fail");
            }
            PyElement::new(dict.into_any().unbind())
        })
    });

    Ok(PyStream(py_stream))
}

/// State for the CSV writer: column order and a line-buffered writer
/// (flushes on each newline so data is visible without waiting for drop).
struct CsvWriteState {
    file: std::io::LineWriter<std::fs::File>,
    headers: Option<Vec<String>>,
}

impl CsvWriteState {
    fn new(path: &str) -> anyhow::Result<Self> {
        let file = std::fs::File::create(path)
            .with_context(|| format!("failed to open CSV file for writing: {path}"))?;
        Ok(Self {
            file: std::io::LineWriter::new(file),
            headers: None,
        })
    }

    fn write_row(&mut self, values: &[String]) {
        let line = values
            .iter()
            .enumerate()
            .fold(String::new(), |mut acc, (i, v)| {
                if i > 0 {
                    acc.push(',');
                }
                // Simple CSV escaping: quote if contains comma, quote, or newline
                if v.contains(',') || v.contains('"') || v.contains('\n') {
                    acc.push('"');
                    acc.push_str(&v.replace('"', "\"\""));
                    acc.push('"');
                } else {
                    acc.push_str(v);
                }
                acc
            });
        if let Err(e) = writeln!(self.file, "{line}") {
            log::error!("csv_write: failed to write row: {e}");
        }
    }
}

/// Inner implementation for the `.csv_write()` stream method.
///
/// Writes each dict as a CSV row. Headers are inferred from the first dict's keys.
/// A `time` column is prepended with the graph time in nanoseconds.
pub fn py_csv_write_inner(
    stream: &Rc<dyn Stream<PyElement>>,
    path: String,
) -> anyhow::Result<Rc<dyn Node>> {
    let state: RefCell<CsvWriteState> = RefCell::new(CsvWriteState::new(&path)?);

    Ok(stream.for_each(move |elem, time| {
        Python::attach(|py| {
            let obj = elem.as_ref().bind(py);
            if let Ok(dict) = obj.cast::<PyDict>() {
                let mut guard = state.borrow_mut();

                if guard.headers.is_none() {
                    let keys: Vec<String> = dict
                        .keys()
                        .iter()
                        .map(|k| k.extract::<String>().unwrap_or_default())
                        .collect();
                    let mut header_row = vec!["time".to_string()];
                    header_row.extend(keys.clone());
                    guard.write_row(&header_row);
                    guard.headers = Some(keys);
                }

                let h = guard
                    .headers
                    .clone()
                    .expect("invariant: headers set above before this point");
                let time_nanos: u64 = time.into();
                let mut values = vec![time_nanos.to_string()];
                for key in &h {
                    let v = dict
                        .get_item(key)
                        .ok()
                        .flatten()
                        .map(|v| v.str().map(|s| s.to_string()).unwrap_or_default())
                        .unwrap_or_default();
                    values.push(v);
                }
                guard.write_row(&values);
            }
        });
    }))
}