Skip to content

Commit

Permalink
wip support json cow
Browse files Browse the repository at this point in the history
  • Loading branch information
davidhewitt committed Mar 18, 2024
1 parent 352d40f commit c26fd87
Show file tree
Hide file tree
Showing 8 changed files with 71 additions and 49 deletions.
3 changes: 1 addition & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Expand Up @@ -44,7 +44,7 @@ base64 = "0.21.7"
num-bigint = "0.4.4"
python3-dll-a = "0.2.7"
uuid = "1.7.0"
jiter = {version = "0.0.7", features = ["python"]}
jiter = { git = "https://github.com/pydantic/jiter", branch = "string-cow", features = ["python"] }

[lib]
name = "_pydantic_core"
Expand Down
2 changes: 1 addition & 1 deletion src/errors/line_error.rs
Expand Up @@ -151,7 +151,7 @@ impl ValLineError {
#[derive(Clone)]
pub enum InputValue {
Python(PyObject),
Json(JsonValue),
Json(JsonValue<'static>),
}

impl ToPyObject for InputValue {
Expand Down
7 changes: 7 additions & 0 deletions src/errors/location.rs
@@ -1,5 +1,6 @@
use pyo3::exceptions::PyTypeError;
use pyo3::sync::GILOnceCell;
use std::borrow::Cow;
use std::fmt;

use pyo3::prelude::*;
Expand Down Expand Up @@ -52,6 +53,12 @@ impl From<&str> for LocItem {
}
}

impl From<Cow<'_, str>> for LocItem {
fn from(s: Cow<'_, str>) -> Self {
Self::S(s.into_owned())
}
}

impl From<i64> for LocItem {
fn from(i: i64) -> Self {
Self::I(i)
Expand Down
18 changes: 9 additions & 9 deletions src/input/input_json.rs
Expand Up @@ -21,26 +21,26 @@ use super::{
};

/// This is required but since JSON object keys are always strings, I don't think it can be called
impl From<&JsonValue> for LocItem {
impl From<&JsonValue<'_>> for LocItem {
fn from(json_value: &JsonValue) -> Self {
match json_value {
JsonValue::Int(i) => (*i).into(),
JsonValue::Str(s) => s.as_str().into(),
JsonValue::Str(s) => s.clone().into(),
v => format!("{v:?}").into(),
}
}
}

impl From<JsonValue> for LocItem {
impl From<JsonValue<'_>> for LocItem {
fn from(json_value: JsonValue) -> Self {
(&json_value).into()
}
}

impl<'py> Input<'py> for JsonValue {
impl<'py> Input<'py> for JsonValue<'_> {
fn as_error_value(&self) -> InputValue {
// cloning JsonValue is cheap due to use of Arc
InputValue::Json(self.clone())
InputValue::Json(self.clone().into_static())
}

fn is_none(&self) -> bool {
Expand Down Expand Up @@ -91,7 +91,7 @@ impl<'py> Input<'py> for JsonValue {
// TODO: in V3 we may want to make JSON str always win if in union, for consistency,
// see https://github.com/pydantic/pydantic-core/pull/867#discussion_r1386582501
match self {
JsonValue::Str(s) => Ok(ValidationMatch::strict(s.as_str().into())),
JsonValue::Str(s) => Ok(ValidationMatch::strict(s.as_ref().into())),
JsonValue::Int(i) if !strict && coerce_numbers_to_str => Ok(ValidationMatch::lax(i.to_string().into())),
JsonValue::BigInt(b) if !strict && coerce_numbers_to_str => Ok(ValidationMatch::lax(b.to_string().into())),
JsonValue::Float(f) if !strict && coerce_numbers_to_str => Ok(ValidationMatch::lax(f.to_string().into())),
Expand Down Expand Up @@ -135,7 +135,7 @@ impl<'py> Input<'py> for JsonValue {

fn exact_str(&self) -> ValResult<EitherString<'_>> {
match self {
JsonValue::Str(s) => Ok(s.as_str().into()),
JsonValue::Str(s) => Ok(s.as_ref().into()),
_ => Err(ValError::new(ErrorTypeDefaults::StringType, self)),
}
}
Expand Down Expand Up @@ -313,7 +313,7 @@ impl<'py> Input<'py> for str {
fn as_error_value(&self) -> InputValue {
// Justification for the clone: this is on the error pathway and we are generally ok
// with errors having a performance penalty
InputValue::Json(JsonValue::Str(self.to_owned()))
InputValue::Json(JsonValue::Str(self.to_owned().into()))
}

fn as_kwargs(&self, _py: Python<'py>) -> Option<Bound<'py, PyDict>> {
Expand Down Expand Up @@ -447,5 +447,5 @@ impl BorrowInput<'_> for String {
}

fn string_to_vec(s: &str) -> JsonArray {
JsonArray::new(s.chars().map(|c| JsonValue::Str(c.to_string())).collect())
JsonArray::new(s.chars().map(|c| JsonValue::Str(c.to_string().into())).collect())
}
68 changes: 42 additions & 26 deletions src/input/return_enums.rs
Expand Up @@ -83,8 +83,8 @@ pub enum GenericIterable<'a, 'py> {
PyByteArray(&'a Bound<'py, PyByteArray>),
Sequence(&'a Bound<'py, PySequence>),
Iterator(Bound<'py, PyIterator>),
JsonArray(&'a [JsonValue]),
JsonObject(&'a JsonObject),
JsonArray(&'a [JsonValue<'a>]),
JsonObject(&'a JsonObject<'a>),
JsonString(&'a str),
}

Expand Down Expand Up @@ -434,7 +434,7 @@ pub enum GenericMapping<'a, 'py> {
PyMapping(&'a Bound<'py, PyMapping>),
StringMapping(&'a Bound<'py, PyDict>),
PyGetAttr(Bound<'py, PyAny>, Option<Bound<'py, PyDict>>),
JsonObject(&'a JsonObject),
JsonObject(&'a JsonObject<'a>),
}

macro_rules! derive_from {
Expand All @@ -450,8 +450,8 @@ macro_rules! derive_from {
derive_from!(GenericMapping, PyDict, PyDict);
derive_from!(GenericMapping, PyMapping, PyMapping);

impl<'a> From<&'a JsonObject> for GenericMapping<'a, '_> {
fn from(s: &'a JsonObject) -> Self {
impl<'a> From<&'a JsonObject<'a>> for GenericMapping<'a, '_> {
fn from(s: &'a JsonObject<'a>) -> Self {
Self::JsonObject(s)
}
}
Expand Down Expand Up @@ -613,41 +613,50 @@ impl<'py> Iterator for AttributesGenericIterator<'py> {
// size_hint is omitted as it isn't needed
}

pub struct JsonObjectGenericIterator<'py> {
object_iter: SliceIter<'py, (String, JsonValue)>,
pub struct JsonObjectGenericIterator<'a> {
object_iter: SliceIter<'a, (Cow<'a, str>, JsonValue<'a>)>,
}

impl<'py> JsonObjectGenericIterator<'py> {
pub fn new(json_object: &'py JsonObject) -> ValResult<Self> {
impl<'a> JsonObjectGenericIterator<'a> {
pub fn new(json_object: &'a JsonObject<'a>) -> ValResult<Self> {
Ok(Self {
object_iter: json_object.iter(),
})
}
}

impl<'py> Iterator for JsonObjectGenericIterator<'py> {
type Item = ValResult<(&'py String, &'py JsonValue)>;
impl<'a> Iterator for JsonObjectGenericIterator<'a> {
type Item = ValResult<(&'a str, &'a JsonValue<'a>)>;

fn next(&mut self) -> Option<Self::Item> {
self.object_iter.next().map(|(key, value)| Ok((key, value)))
self.object_iter.next().map(|(key, value)| Ok((key.as_ref(), value)))
}
// size_hint is omitted as it isn't needed
}

#[derive(Debug, Clone)]
pub enum GenericIterator {
#[derive(Debug)]
pub enum GenericIterator<'a> {
PyIterator(GenericPyIterator),
JsonArray(GenericJsonIterator),
JsonArray(GenericJsonIterator<'a>),
}

impl From<JsonArray> for GenericIterator {
fn from(array: JsonArray) -> Self {
impl GenericIterator<'_> {
pub(crate) fn into_static(self) -> GenericIterator<'static> {
match self {
GenericIterator::PyIterator(iter) => GenericIterator::PyIterator(iter),
GenericIterator::JsonArray(iter) => GenericIterator::JsonArray(iter.into_static()),
}
}
}

impl<'a> From<JsonArray<'a>> for GenericIterator<'a> {
fn from(array: JsonArray<'a>) -> Self {
let json_iter = GenericJsonIterator { array, index: 0 };
Self::JsonArray(json_iter)
}
}

impl From<&Bound<'_, PyAny>> for GenericIterator {
impl From<&Bound<'_, PyAny>> for GenericIterator<'_> {
fn from(obj: &Bound<'_, PyAny>) -> Self {
let py_iter = GenericPyIterator {
obj: obj.clone().into(),
Expand Down Expand Up @@ -690,13 +699,13 @@ impl GenericPyIterator {
}

#[derive(Debug, Clone)]
pub struct GenericJsonIterator {
array: JsonArray,
pub struct GenericJsonIterator<'a> {
array: JsonArray<'a>,
index: usize,
}

impl GenericJsonIterator {
pub fn next(&mut self, _py: Python) -> PyResult<Option<(&JsonValue, usize)>> {
impl<'a> GenericJsonIterator<'a> {
pub fn next(&mut self, _py: Python) -> PyResult<Option<(&JsonValue<'a>, usize)>> {
if self.index < self.array.len() {
// panic here is impossible due to bounds check above; compiler should be
// able to optimize it away even
Expand All @@ -710,12 +719,19 @@ impl GenericJsonIterator {
}

pub fn input_as_error_value(&self, _py: Python<'_>) -> InputValue {
InputValue::Json(JsonValue::Array(self.array.clone()))
InputValue::Json(JsonValue::Array(self.array.clone()).into_static())
}

pub fn index(&self) -> usize {
self.index
}

pub fn into_static(self) -> GenericJsonIterator<'static> {
GenericJsonIterator {
array: JsonArray::new(self.array.iter().map(|v| v.clone().into_static()).collect()),
index: self.index,
}
}
}

#[cfg_attr(debug_assertions, derive(Debug))]
Expand All @@ -732,12 +748,12 @@ impl<'py> PyArgs<'py> {

#[cfg_attr(debug_assertions, derive(Debug))]
pub struct JsonArgs<'a> {
pub args: Option<&'a [JsonValue]>,
pub kwargs: Option<&'a JsonObject>,
pub args: Option<&'a [JsonValue<'a>]>,
pub kwargs: Option<&'a JsonObject<'a>>,
}

impl<'a> JsonArgs<'a> {
pub fn new(args: Option<&'a [JsonValue]>, kwargs: Option<&'a JsonObject>) -> Self {
pub fn new(args: Option<&'a [JsonValue<'a>]>, kwargs: Option<&'a JsonObject<'a>>) -> Self {
Self { args, kwargs }
}
}
Expand Down
16 changes: 8 additions & 8 deletions src/lookup_key.rs
Expand Up @@ -262,10 +262,10 @@ impl LookupKey {

pub fn json_get<'data, 's>(
&'s self,
dict: &'data JsonObject,
) -> ValResult<Option<(&'s LookupPath, &'data JsonValue)>> {
dict: &'data JsonObject<'data>,
) -> ValResult<Option<(&'s LookupPath, &'data JsonValue<'data>)>> {
match self {
Self::Simple { key, path, .. } => match dict.get(key) {
Self::Simple { key, path, .. } => match dict.get(key.as_str()) {
Some(value) => Ok(Some((path, value))),
None => Ok(None),
},
Expand All @@ -275,9 +275,9 @@ impl LookupKey {
key2,
path2,
..
} => match dict.get(key1) {
} => match dict.get(key1.as_str()) {
Some(value) => Ok(Some((path1, value))),
None => match dict.get(key2) {
None => match dict.get(key2.as_str()) {
Some(value) => Ok(Some((path2, value))),
None => Ok(None),
},
Expand Down Expand Up @@ -475,7 +475,7 @@ impl PathItem {
}
}

pub fn json_get<'a>(&self, any_json: &'a JsonValue) -> Option<&'a JsonValue> {
pub fn json_get<'a>(&self, any_json: &'a JsonValue<'a>) -> Option<&'a JsonValue<'a>> {
match any_json {
JsonValue::Object(v_obj) => self.json_obj_get(v_obj),
JsonValue::Array(v_array) => match self {
Expand All @@ -493,9 +493,9 @@ impl PathItem {
}
}

pub fn json_obj_get<'a>(&self, json_obj: &'a JsonObject) -> Option<&'a JsonValue> {
pub fn json_obj_get<'a>(&self, json_obj: &'a JsonObject<'a>) -> Option<&'a JsonValue<'a>> {
match self {
Self::S(key, _) => json_obj.get(key),
Self::S(key, _) => json_obj.get(key.as_str()),
_ => None,
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/validators/generator.rs
Expand Up @@ -65,7 +65,7 @@ impl Validator for GeneratorValidator {
input: &(impl Input<'py> + ?Sized),
state: &mut ValidationState<'_, 'py>,
) -> ValResult<PyObject> {
let iterator = input.validate_iter()?;
let iterator = input.validate_iter()?.into_static();
let validator = self.item_validator.as_ref().map(|v| {
InternalValidator::new(
py,
Expand Down Expand Up @@ -96,7 +96,7 @@ impl Validator for GeneratorValidator {
#[pyclass(module = "pydantic_core._pydantic_core")]
#[derive(Debug)]
struct ValidatorIterator {
iterator: GenericIterator,
iterator: GenericIterator<'static>,
validator: Option<InternalValidator>,
min_length: Option<usize>,
max_length: Option<usize>,
Expand Down

0 comments on commit c26fd87

Please sign in to comment.