Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(python): Add conversion helper functions to Python extension. #231

Merged
merged 3 commits into from
Aug 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pywr-python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ pyo3-log = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
chrono = { workspace = true }

pywr-v1-schema = { workspace = true }
pywr-core = { path = "../pywr-core" }
pywr-schema = { path = "../pywr-schema" }

Expand Down
2 changes: 1 addition & 1 deletion pywr-python/pywr/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from pathlib import Path
from typing import Optional

from .pywr import Schema, Model # type: ignore
from .pywr import Schema, Model, convert_model_from_v1_json_string, convert_metric_from_v1_json_string # type: ignore


def run_from_path(
Expand Down
52 changes: 50 additions & 2 deletions pywr-python/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use chrono::NaiveDateTime;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyType};
use pyo3::types::{PyDict, PyTuple, PyType};

/// Python API
///
Expand All @@ -16,6 +16,8 @@ use pywr_core::solvers::{ClpSolver, ClpSolverSettings, ClpSolverSettingsBuilder}
#[cfg(feature = "highs")]
use pywr_core::solvers::{HighsSolver, HighsSolverSettings, HighsSolverSettingsBuilder};
use pywr_schema::model::DateType;
use pywr_schema::parameters::TryIntoV2Parameter;
use pywr_schema::ConversionError;
use std::fmt;
use std::path::PathBuf;
use std::str::FromStr;
Expand Down Expand Up @@ -95,13 +97,56 @@ impl Schema {
Ok(data)
}

/// Convert the schema to a Pywr model.
/// Build the schema in to a Pywr model.
fn build(&mut self, data_path: Option<PathBuf>, output_path: Option<PathBuf>) -> PyResult<Model> {
let model = self.schema.build_model(data_path.as_deref(), output_path.as_deref())?;
Ok(Model { model })
}
}

/// Convert a Pywr v1.x JSON string to a Pywr v2.x schema.
#[pyfunction]
fn convert_model_from_v1_json_string(py: Python, data: &str) -> PyResult<Py<PyTuple>> {
// Try to convert
let (schema, errors) =
pywr_schema::PywrModel::from_v1_str(data).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

// Create a new schema object
let py_schema = Schema { schema };
let py_errors = errors.into_iter().map(|e| e.to_string()).collect::<Vec<_>>();

let result = PyTuple::new_bound(py, &[py_schema.into_py(py), py_errors.into_py(py)]).into();
Ok(result)
}

#[pyclass]
pub struct Metric {
metric: pywr_schema::metric::Metric,
}

#[pymethods]
impl Metric {
/// Serialize the metric to a JSON string.
fn to_json_string(&self) -> PyResult<String> {
let data = serde_json::to_string_pretty(&self.metric).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
Ok(data)
}
}

/// Convert a Pywr v1.x JSON string to a Pywr v2.x metric.
#[pyfunction]
fn convert_metric_from_v1_json_string(_py: Python, data: &str) -> PyResult<Metric> {
let v1: pywr_v1_schema::parameters::ParameterValue =
serde_json::from_str(data).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;

let metric = v1
.try_into_v2_parameter(None, &mut 0)
.map_err(|e: ConversionError| PyRuntimeError::new_err(e.to_string()))?;

let py_metric = Metric { metric };
Ok(py_metric)
}

#[pyclass]
pub struct Model {
model: pywr_core::models::Model,
Expand Down Expand Up @@ -203,8 +248,11 @@ fn build_highs_settings(kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<HighsSol
fn pywr(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
pyo3_log::init();

m.add_function(wrap_pyfunction!(convert_model_from_v1_json_string, m)?)?;
m.add_function(wrap_pyfunction!(convert_metric_from_v1_json_string, m)?)?;
m.add_class::<Schema>()?;
m.add_class::<Model>()?;
m.add_class::<Metric>()?;

Ok(())
}
9 changes: 9 additions & 0 deletions pywr-schema/src/model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -667,6 +667,15 @@ impl PywrModel {
errors,
)
}

/// Convert a v1 JSON string to a v2 model.
///
/// See [`PywrModel::from_v1`] for more information.
pub fn from_v1_str(v1: &str) -> Result<(Self, Vec<ConversionError>), pywr_v1_schema::model::PywrSchemaError> {
let v1_model: pywr_v1_schema::PywrModel = serde_json::from_str(v1)?;

Ok(Self::from_v1(v1_model))
}
}

#[derive(serde::Deserialize, serde::Serialize, Clone)]
Expand Down