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.

Sign up for GitHub

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

Jump to bottom

Support for parsing parial JSON strings in Python #66

Merged
merged 3 commits into from Mar 25, 2024
Merged

Support for parsing parial JSON strings in Python #66

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ed71053
add "allow_partial" argument to python parsing
samuelcolvin Feb 6, 2024
c8a669e
Merge branch 'main' into partial-parsing
samuelcolvin Mar 25, 2024
c3f6b6b
separate _parse_array and _parse_object
samuelcolvin Mar 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 3 additions & 2 deletions benches/python.rs
View file
Open in desktop
Expand Up @@ -15,6 +15,7 @@ fn python_parse_numeric(bench: &mut Bencher) {
br#" { "int": 1, "bigint": 123456789012345678901234567890, "float": 1.2} "#,
false,
true,
false,
)
.unwrap()
});
Expand All @@ -23,7 +24,7 @@ fn python_parse_numeric(bench: &mut Bencher) {

fn python_parse_other(bench: &mut Bencher) {
Python::with_gil(|py| {
bench.iter(|| python_parse(py, br#"["string", true, false, null]"#, false, true).unwrap());
bench.iter(|| python_parse(py, br#"["string", true, false, null]"#, false, true, false).unwrap());
})
}

Expand All @@ -34,7 +35,7 @@ fn _python_parse_file(path: &str, bench: &mut Bencher, cache_strings: bool) {
let json_data = contents.as_bytes();

Python::with_gil(|py| {
bench.iter(|| python_parse(py, json_data, false, cache_strings).unwrap());
bench.iter(|| python_parse(py, json_data, false, cache_strings, false).unwrap());
})
}

Expand Down
118 changes: 90 additions & 28 deletions src/python.rs
View file
Open in desktop
Expand Up @@ -13,6 +13,7 @@
use crate::number_decoder::{NumberAny, NumberInt};
use crate::parse::{Parser, Peek};
use crate::string_decoder::{StringDecoder, Tape};
use crate::JsonErrorType;

/// Parse a JSON value from a byte slice and return a Python object.
///
Expand All @@ -31,12 +32,14 @@
json_data: &[u8],
allow_inf_nan: bool,
cache_strings: bool,
allow_partial: bool,
) -> JsonResult<Bound<'py, PyAny>> {
let mut python_parser = PythonParser {
parser: Parser::new(json_data),
tape: Tape::default(),
recursion_limit: DEFAULT_RECURSION_LIMIT,
allow_inf_nan,
allow_partial,
};

let peek = python_parser.parser.peek()?;
Expand All @@ -45,7 +48,9 @@
} else {
python_parser.py_take_value::<StringNoCache>(py, peek)?
};
python_parser.parser.finish()?;
if !allow_partial {
python_parser.parser.finish()?;
}
Ok(v)
}

Expand All @@ -59,6 +64,7 @@
tape: Tape,
recursion_limit: u8,
allow_inf_nan: bool,
allow_partial: bool,
}

impl<'j> PythonParser<'j> {
Expand All @@ -67,6 +73,21 @@
py: Python<'py>,
peek: Peek,
) -> JsonResult<Bound<'py, PyAny>> {
macro_rules! tri {
($result:expr, $partial_value:expr) => {
match $result {
Ok(k) => k,
Err(e) => {
return if self._allow_partial_err(&e) {
Ok($partial_value.into_any())
} else {
Err(e)
}
}
}
};
}

match peek {
Peek::Null => {
self.parser.consume_null()?;
Expand All @@ -85,13 +106,12 @@
Ok(StringCache::get(py, s.as_str()))
}
Peek::Array => {
let list = if let Some(peek_first) = self.parser.array_first()? {
let list = if let Some(peek_first) = tri!(self.parser.array_first(), PyList::empty_bound(py)) {
let mut vec: SmallVec<[Bound<'_, PyAny>; 8]> = SmallVec::with_capacity(8);
let v = self._check_take_value::<StringCache>(py, peek_first)?;
vec.push(v);
while let Some(peek) = self.parser.array_step()? {
let v = self._check_take_value::<StringCache>(py, peek)?;
vec.push(v);
if let Err(e) = self._parse_array::<StringCache>(py, peek_first, &mut vec) {
if !self._allow_partial_err(&e) {
return Err(e);
}
}
PyList::new_bound(py, vec)
} else {
Expand All @@ -101,27 +121,9 @@
}
Peek::Object => {
let dict = PyDict::new_bound(py);

let set_item = |key: Bound<'py, PyAny>, value: Bound<'py, PyAny>| {
let r = unsafe { ffi::PyDict_SetItem(dict.as_ptr(), key.as_ptr(), value.as_ptr()) };
// AFAIK this shouldn't happen since the key will always be a string which is hashable
// we panic here rather than returning a result and using `?` below as it's up to 14% faster
// presumably because there are fewer branches
if r == -1 {
panic!("PyDict_SetItem failed")
}
};

if let Some(first_key) = self.parser.object_first::<StringDecoder>(&mut self.tape)? {
let first_key = StringCache::get(py, first_key.as_str());
let peek = self.parser.peek()?;
let first_value = self._check_take_value::<StringCache>(py, peek)?;
set_item(first_key, first_value);
while let Some(key) = self.parser.object_step::<StringDecoder>(&mut self.tape)? {
let key = StringCache::get(py, key.as_str());
let peek = self.parser.peek()?;
let value = self._check_take_value::<StringCache>(py, peek)?;
set_item(key, value);
if let Err(e) = self._parse_object::<StringCache>(py, &dict) {
if !self._allow_partial_err(&e) {
return Err(e);

Check warning on line 126 in src/python.rs

View check run for this annotation

Codecov / codecov/patch

src/python.rs#L126 

Added line #L126 was not covered by tests
}
}
Ok(dict.into_any())
Expand All @@ -146,6 +148,66 @@
}
}

fn _parse_array<'py, StringCache: StringMaybeCache>(
&mut self,
py: Python<'py>,
peek_first: Peek,
vec: &mut SmallVec<[Bound<'py, PyAny>; 8]>,
) -> JsonResult<()> {
let v = self._check_take_value::<StringCache>(py, peek_first)?;
vec.push(v);
while let Some(peek) = self.parser.array_step()? {
let v = self._check_take_value::<StringCache>(py, peek)?;
vec.push(v);
}
Ok(())
}

fn _parse_object<'py, StringCache: StringMaybeCache>(
&mut self,
py: Python<'py>,
dict: &Bound<'py, PyDict>,
) -> JsonResult<()> {
let set_item = |key: Bound<'py, PyAny>, value: Bound<'py, PyAny>| {
let r = unsafe { ffi::PyDict_SetItem(dict.as_ptr(), key.as_ptr(), value.as_ptr()) };
// AFAIK this shouldn't happen since the key will always be a string which is hashable
// we panic here rather than returning a result and using `?` below as it's up to 14% faster
// presumably because there are fewer branches
if r == -1 {
panic!("PyDict_SetItem failed")

Check warning on line 177 in src/python.rs

View check run for this annotation

Codecov / codecov/patch

src/python.rs#L177 

Added line #L177 was not covered by tests
}
};
if let Some(first_key) = self.parser.object_first::<StringDecoder>(&mut self.tape)? {
let first_key = StringCache::get(py, first_key.as_str());
let peek = self.parser.peek()?;
let first_value = self._check_take_value::<StringCache>(py, peek)?;
set_item(first_key, first_value);
while let Some(key) = self.parser.object_step::<StringDecoder>(&mut self.tape)? {
let key = StringCache::get(py, key.as_str());
let peek = self.parser.peek()?;
let value = self._check_take_value::<StringCache>(py, peek)?;
set_item(key, value);
}
}

Check warning on line 191 in src/python.rs

View check run for this annotation

Codecov / codecov/patch

src/python.rs#L191 

Added line #L191 was not covered by tests
Ok(())
}

fn _allow_partial_err(&self, e: &JsonError) -> bool {
if self.allow_partial {
matches!(

Check warning on line 197 in src/python.rs

View check run for this annotation

Codecov / codecov/patch

src/python.rs#L197 

Added line #L197 was not covered by tests
e.error_type,
JsonErrorType::EofWhileParsingList
| JsonErrorType::EofWhileParsingObject
| JsonErrorType::EofWhileParsingString
| JsonErrorType::EofWhileParsingValue
| JsonErrorType::ExpectedListCommaOrEnd
| JsonErrorType::ExpectedObjectCommaOrEnd
)
} else {
false
}
}

fn _check_take_value<'py, StringCache: StringMaybeCache>(
&mut self,
py: Python<'py>,
Expand Down
67 changes: 59 additions & 8 deletions tests/python.rs
View file
Open in desktop
@@ -1,4 +1,5 @@
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use pyo3::ToPyObject;

use jiter::{map_json_error, python_parse, JsonValue};
Expand Down Expand Up @@ -42,6 +43,7 @@ fn test_python_parse_numeric() {
br#" { "int": 1, "bigint": 123456789012345678901234567890, "float": 1.2} "#,
false,
true,
false,
)
.unwrap();
assert_eq!(
Expand All @@ -59,6 +61,7 @@ fn test_python_parse_other_cached() {
br#"["string", true, false, null, NaN, Infinity, -Infinity]"#,
true,
true,
false,
)
.unwrap();
assert_eq!(obj.to_string(), "['string', True, False, None, nan, inf, -inf]");
Expand All @@ -68,15 +71,15 @@ fn test_python_parse_other_cached() {
#[test]
fn test_python_parse_other_no_cache() {
Python::with_gil(|py| {
let obj = python_parse(py, br#"["string", true, false, null]"#, false, false).unwrap();
let obj = python_parse(py, br#"["string", true, false, null]"#, false, false, false).unwrap();
assert_eq!(obj.to_string(), "['string', True, False, None]");
})
}

#[test]
fn test_python_disallow_nan() {
Python::with_gil(|py| {
let r = python_parse(py, br#"[NaN]"#, false, true);
let r = python_parse(py, br#"[NaN]"#, false, true, false);
let e = r.map_err(|e| map_json_error(br#"[NaN]"#, &e)).unwrap_err();
assert_eq!(e.to_string(), "ValueError: expected value at line 1 column 2");
})
Expand All @@ -86,7 +89,7 @@ fn test_python_disallow_nan() {
fn test_error() {
Python::with_gil(|py| {
let bytes = br#"["string""#;
let r = python_parse(py, bytes, false, true);
let r = python_parse(py, bytes, false, true, false);
let e = r.map_err(|e| map_json_error(bytes, &e)).unwrap_err();
assert_eq!(e.to_string(), "ValueError: EOF while parsing a list at line 1 column 9");
})
Expand All @@ -98,7 +101,7 @@ fn test_recursion_limit() {
let bytes = json.as_bytes();

Python::with_gil(|py| {
let r = python_parse(py, bytes, false, true);
let r = python_parse(py, bytes, false, true, false);
let e = r.map_err(|e| map_json_error(bytes, &e)).unwrap_err();
assert_eq!(
e.to_string(),
Expand All @@ -114,24 +117,72 @@ fn test_recursion_limit_incr() {
let bytes = json.as_bytes();

Python::with_gil(|py| {
let v = python_parse(py, bytes, false, true).unwrap();
let v = python_parse(py, bytes, false, true, false).unwrap();
assert_eq!(v.len().unwrap(), 2000);
});

Python::with_gil(|py| {
let v = python_parse(py, bytes, false, true).unwrap();
let v = python_parse(py, bytes, false, true, false).unwrap();
assert_eq!(v.len().unwrap(), 2000);
});
}

#[test]
fn test_exected_value_error() {
fn test_extracted_value_error() {
let json = "xx";
let bytes = json.as_bytes();

Python::with_gil(|py| {
let r = python_parse(py, bytes, false, true);
let r = python_parse(py, bytes, false, true, false);
let e = r.map_err(|e| map_json_error(bytes, &e)).unwrap_err();
assert_eq!(e.to_string(), "ValueError: expected value at line 1 column 1");
})
}

#[test]
fn test_partial_array() {
Python::with_gil(|py| {
let bytes = br#"["string", true, null, 1, "foo"#;
let py_value = python_parse(py, bytes, false, true, true).unwrap();
let string = py_value.to_string();
assert_eq!(string, "['string', True, None, 1]");

// test that stopping at every points is ok
for i in 1..bytes.len() {
let py_value = python_parse(py, &bytes[..i], false, true, true).unwrap();
assert!(py_value.is_instance_of::<PyList>());
}
})
}

#[test]
fn test_partial_object() {
Python::with_gil(|py| {
let bytes = br#"{"a": 1, "b": 2, "c"#;
let py_value = python_parse(py, bytes, false, true, true).unwrap();
let string = py_value.to_string();
assert_eq!(string, "{'a': 1, 'b': 2}");

// test that stopping at every points is ok
for i in 1..bytes.len() {
let py_value = python_parse(py, &bytes[..i], false, true, true).unwrap();
assert!(py_value.is_instance_of::<PyDict>());
}
})
}

#[test]
fn test_partial_nested() {
Python::with_gil(|py| {
let bytes = br#"{"a": 1, "b": 2, "c": [1, 2, {"d": 1, "#;
let py_value = python_parse(py, bytes, false, true, true).unwrap();
let string = py_value.to_string();
assert_eq!(string, "{'a': 1, 'b': 2, 'c': [1, 2, {'d': 1}]}");

// test that stopping at every points is ok
for i in 1..bytes.len() {
let py_value = python_parse(py, &bytes[..i], false, true, true).unwrap();
assert!(py_value.is_instance_of::<PyDict>());
}
})
}