Consistent use of PyString and to_str (#210)

samuelcolvin · web-flow · commit d30d268584f2 · 2022-08-01T17:59:08.000+01:00
* Consistent use of PyString and to_str

* tests for invalid unicode containing unpaired surrogates
diff --git a/src/errors/line_error.rs b/src/errors/line_error.rs
@@ -43,6 +43,10 @@ impl<'a> ValError<'a> {
         Self::LineErrors(vec![ValLineError::new_with_loc(kind, input, loc)])
     }
 
+    pub fn new_custom_input(kind: ErrorKind, input_value: InputValue<'a>) -> ValError<'a> {
+        Self::LineErrors(vec![ValLineError::new_custom_input(kind, input_value)])
+    }
+
     /// helper function to call with_outer on line items if applicable
     pub fn with_outer_location(self, loc_item: LocItem) -> Self {
         match self {
@@ -90,6 +94,14 @@ impl<'a> ValLineError<'a> {
         }
     }
 
+    pub fn new_custom_input(kind: ErrorKind, input_value: InputValue<'a>) -> ValLineError<'a> {
+        Self {
+            kind,
+            input_value,
+            location: Location::default(),
+        }
+    }
+
     /// location is stored reversed so it's quicker to add "outer" items as that's what we always do
     /// hence `push` here instead of `insert`
     pub fn with_outer_location(mut self, loc_item: LocItem) -> Self {
diff --git a/src/errors/value_exception.rs b/src/errors/value_exception.rs
@@ -46,8 +46,8 @@ impl PydanticValueError {
         if let Some(ref context) = self.context {
             for item in context.as_ref(py).items().iter() {
                 let (key, value): (&PyString, &PyAny) = item.extract()?;
-                if let Ok(value_str) = value.extract::<&PyString>() {
-                    message = message.replace(&format!("{{{}}}", key.to_str()?), value_str.to_str()?);
+                if let Ok(py_str) = value.cast_as::<PyString>() {
+                    message = message.replace(&format!("{{{}}}", key.to_str()?), py_str.to_str()?);
                 } else if let Ok(value_int) = value.extract::<i64>() {
                     message = message.replace(&format!("{{{}}}", key.to_str()?), &value_int.to_string());
                 } else {
diff --git a/src/input/input_python.rs b/src/input/input_python.rs
@@ -18,14 +18,14 @@ use super::datetime::{
 };
 use super::shared::{float_as_int, int_as_bool, str_as_bool, str_as_int};
 use super::{
-    repr_string, EitherBytes, EitherString, EitherTimedelta, GenericArguments, GenericListLike, GenericMapping, Input,
-    PyArgs,
+    py_string_str, repr_string, EitherBytes, EitherString, EitherTimedelta, GenericArguments, GenericListLike,
+    GenericMapping, Input, PyArgs,
 };
 
 impl<'a> Input<'a> for PyAny {
     fn as_loc_item(&self) -> LocItem {
-        if let Ok(key_str) = self.extract::<String>() {
-            key_str.into()
+        if let Ok(py_str) = self.cast_as::<PyString>() {
+            py_str.to_string_lossy().as_ref().into()
         } else if let Ok(key_int) = self.extract::<usize>() {
             key_int.into()
         } else {
@@ -140,8 +140,8 @@ impl<'a> Input<'a> for PyAny {
         if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
             Ok(py_bytes.into())
         } else if let Ok(py_str) = self.cast_as::<PyString>() {
-            let string = py_str.to_string_lossy().to_string();
-            Ok(string.into_bytes().into())
+            let str = py_string_str(py_str)?;
+            Ok(str.as_bytes().into())
         } else if let Ok(py_byte_array) = self.cast_as::<PyByteArray>() {
             Ok(py_byte_array.to_vec().into())
         } else {
@@ -369,7 +369,8 @@ impl<'a> Input<'a> for PyAny {
             Err(ValError::new(ErrorKind::DateType, self))
         } else if let Ok(date) = self.cast_as::<PyDate>() {
             Ok(date.into())
-        } else if let Ok(str) = self.extract::<String>() {
+        } else if let Ok(py_str) = self.cast_as::<PyString>() {
+            let str = py_string_str(py_str)?;
             bytes_as_date(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
             bytes_as_date(self, py_bytes.as_bytes())
@@ -389,7 +390,8 @@ impl<'a> Input<'a> for PyAny {
     fn lax_time(&self) -> ValResult<EitherTime> {
         if let Ok(time) = self.cast_as::<PyTime>() {
             Ok(time.into())
-        } else if let Ok(str) = self.extract::<String>() {
+        } else if let Ok(py_str) = self.cast_as::<PyString>() {
+            let str = py_string_str(py_str)?;
             bytes_as_time(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
             bytes_as_time(self, py_bytes.as_bytes())
@@ -415,7 +417,8 @@ impl<'a> Input<'a> for PyAny {
     fn lax_datetime(&self) -> ValResult<EitherDateTime> {
         if let Ok(dt) = self.cast_as::<PyDateTime>() {
             Ok(dt.into())
-        } else if let Ok(str) = self.extract::<String>() {
+        } else if let Ok(py_str) = self.cast_as::<PyString>() {
+            let str = py_string_str(py_str)?;
             bytes_as_datetime(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
             bytes_as_datetime(self, py_bytes.as_bytes())
@@ -444,7 +447,8 @@ impl<'a> Input<'a> for PyAny {
         if let Ok(dt) = self.cast_as::<PyDelta>() {
             Ok(dt.into())
         } else if let Ok(py_str) = self.cast_as::<PyString>() {
-            bytes_as_timedelta(self, py_str.to_string_lossy().as_bytes())
+            let str = py_string_str(py_str)?;
+            bytes_as_timedelta(self, str.as_bytes())
         } else if let Ok(py_bytes) = self.cast_as::<PyBytes>() {
             bytes_as_timedelta(self, py_bytes.as_bytes())
         } else if let Ok(int) = self.extract::<i64>() {
@@ -518,7 +522,8 @@ fn from_attributes_applicable(obj: &PyAny) -> bool {
 /// Utility for extracting a string from a PyAny, if possible.
 fn maybe_as_string(v: &PyAny, unicode_error: ErrorKind) -> ValResult<Option<Cow<str>>> {
     if let Ok(py_string) = v.cast_as::<PyString>() {
-        Ok(Some(py_string.to_string_lossy()))
+        let str = py_string_str(py_string)?;
+        Ok(Some(Cow::Borrowed(str)))
     } else if let Ok(bytes) = v.cast_as::<PyBytes>() {
         match from_utf8(bytes.as_bytes()) {
             Ok(s) => Ok(Some(Cow::Owned(s.to_string()))),
diff --git a/src/input/mod.rs b/src/input/mod.rs
@@ -12,7 +12,7 @@ pub use datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
 pub use input_abstract::Input;
 pub use parse_json::{JsonInput, JsonObject};
 pub use return_enums::{
-    EitherBytes, EitherString, GenericArguments, GenericListLike, GenericMapping, JsonArgs, PyArgs,
+    py_string_str, EitherBytes, EitherString, GenericArguments, GenericListLike, GenericMapping, JsonArgs, PyArgs,
 };
 
 pub fn repr_string(v: &PyAny) -> PyResult<String> {
diff --git a/src/input/return_enums.rs b/src/input/return_enums.rs
@@ -3,7 +3,7 @@ use std::borrow::Cow;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple};
 
-use crate::errors::{ErrorKind, ValError, ValLineError, ValResult};
+use crate::errors::{ErrorKind, InputValue, ValError, ValLineError, ValResult};
 use crate::recursion_guard::RecursionGuard;
 use crate::validators::{CombinedValidator, Extra, Validator};
 
@@ -212,10 +212,10 @@ pub enum EitherString<'a> {
 }
 
 impl<'a> EitherString<'a> {
-    pub fn as_cow(&self) -> Cow<str> {
+    pub fn as_cow(&self) -> ValResult<'a, Cow<str>> {
         match self {
-            Self::Cow(data) => data.clone(),
-            Self::Py(py_str) => py_str.to_string_lossy(),
+            Self::Cow(data) => Ok(data.clone()),
+            Self::Py(py_str) => Ok(Cow::Borrowed(py_string_str(py_str)?)),
         }
     }
 
@@ -251,6 +251,12 @@ impl<'a> IntoPy<PyObject> for EitherString<'a> {
     }
 }
 
+pub fn py_string_str(py_str: &PyString) -> ValResult<&str> {
+    py_str.to_str().map_err(|_| {
+        ValError::new_custom_input(ErrorKind::StrUnicode, InputValue::PyObject(py_str.into_py(py_str.py())))
+    })
+}
+
 #[cfg_attr(debug_assertions, derive(Debug))]
 pub enum EitherBytes<'a> {
     Cow(Cow<'a, [u8]>),
diff --git a/src/lookup_key.rs b/src/lookup_key.rs
@@ -64,11 +64,9 @@ impl LookupKey {
             };
             let mut locs: Vec<Path> = if first.cast_as::<PyString>().is_ok() {
                 // list of strings rather than list of lists
-                vec![Self::path_choice(py, list)?]
+                vec![Self::path_choice(list)?]
             } else {
-                list.iter()
-                    .map(|obj| Self::path_choice(py, obj))
-                    .collect::<PyResult<_>>()?
+                list.iter().map(Self::path_choice).collect::<PyResult<_>>()?
             };
 
             if let Some(alt_alias) = alt_alias {
@@ -82,12 +80,12 @@ impl LookupKey {
         LookupKey::Simple(key.to_string(), py_string!(py, key))
     }
 
-    fn path_choice(py: Python, obj: &PyAny) -> PyResult<Path> {
+    fn path_choice(obj: &PyAny) -> PyResult<Path> {
         let path = obj
             .extract::<&PyList>()?
             .iter()
             .enumerate()
-            .map(|(index, obj)| PathItem::from_py(py, index, obj))
+            .map(|(index, obj)| PathItem::from_py(index, obj))
             .collect::<PyResult<Path>>()?;
 
         if path.is_empty() {
@@ -237,10 +235,10 @@ fn path_to_string(path: &Path) -> String {
 }
 
 impl PathItem {
-    pub fn from_py(py: Python, index: usize, obj: &PyAny) -> PyResult<Self> {
-        if let Ok(str_key) = obj.extract::<String>() {
-            let py_str_key = py_string!(py, &str_key);
-            Ok(Self::S(str_key, py_str_key))
+    pub fn from_py(index: usize, obj: &PyAny) -> PyResult<Self> {
+        if let Ok(py_str_key) = obj.cast_as::<PyString>() {
+            let str_key = py_str_key.to_str()?.to_string();
+            Ok(Self::S(str_key, py_str_key.into()))
         } else {
             let int_key = obj.extract::<usize>()?;
             if index == 0 {
diff --git a/src/validators/arguments.rs b/src/validators/arguments.rs
@@ -281,7 +281,7 @@ impl Validator for ArgumentsValidator {
                             }
                             Err(err) => return Err(err),
                         };
-                        if !used_kwargs.contains(either_str.as_cow().as_ref()) {
+                        if !used_kwargs.contains(either_str.as_cow()?.as_ref()) {
                             match self.var_kwargs_validator {
                                 Some(ref validator) => match validator.validate(py, value, extra, slots, recursion_guard) {
                                     Ok(value) => output_kwargs.set_item(either_str.as_py_string(py), value)?,
diff --git a/src/validators/literal.rs b/src/validators/literal.rs
@@ -2,7 +2,7 @@ use std::hash::BuildHasherDefault;
 
 use pyo3::intern;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList};
+use pyo3::types::{PyDict, PyList, PyString};
 
 use ahash::AHashSet;
 
@@ -29,8 +29,8 @@ impl BuildValidator for LiteralBuilder {
             return py_error!(r#""expected" should have length > 0"#);
         } else if expected.len() == 1 {
             let first = expected.get_item(0)?;
-            if let Ok(str) = first.extract::<String>() {
-                return Ok(LiteralSingleStringValidator::new(str).into());
+            if let Ok(py_str) = first.cast_as::<PyString>() {
+                return Ok(LiteralSingleStringValidator::new(py_str.to_str()?.to_string()).into());
             }
             if let Ok(int) = first.extract::<i64>() {
                 return Ok(LiteralSingleIntValidator::new(int).into());
@@ -72,7 +72,7 @@ impl Validator for LiteralSingleStringValidator {
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
         let either_str = input.strict_str()?;
-        if either_str.as_cow().as_ref() == self.expected.as_str() {
+        if either_str.as_cow()?.as_ref() == self.expected.as_str() {
             Ok(input.to_object(py))
         } else {
             Err(ValError::new(
@@ -166,7 +166,7 @@ impl Validator for LiteralMultipleStringsValidator {
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
         let either_str = input.strict_str()?;
-        if self.expected.contains(either_str.as_cow().as_ref()) {
+        if self.expected.contains(either_str.as_cow()?.as_ref()) {
             Ok(input.to_object(py))
         } else {
             Err(ValError::new(
@@ -255,8 +255,8 @@ impl LiteralGeneralValidator {
             repr_args.push(item.repr()?.extract()?);
             if let Ok(int) = item.extract::<i64>() {
                 expected_int.insert(int);
-            } else if let Ok(str) = item.extract::<String>() {
-                expected_str.insert(str);
+            } else if let Ok(py_str) = item.cast_as::<PyString>() {
+                expected_str.insert(py_str.to_str()?.to_string());
             } else {
                 expected_py.append(item)?;
             }
@@ -291,7 +291,7 @@ impl Validator for LiteralGeneralValidator {
         }
         if !self.expected_str.is_empty() {
             if let Ok(either_str) = input.strict_str() {
-                if self.expected_str.contains(either_str.as_cow().as_ref()) {
+                if self.expected_str.contains(either_str.as_cow()?.as_ref()) {
                     return Ok(input.to_object(py));
                 }
             }
diff --git a/src/validators/string.rs b/src/validators/string.rs
@@ -90,7 +90,7 @@ impl Validator for StrConstrainedValidator {
         _recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
         let either_str = input.validate_str(extra.strict.unwrap_or(self.strict))?;
-        let cow = either_str.as_cow();
+        let cow = either_str.as_cow()?;
         let mut str = cow.as_ref();
         if let Some(min_length) = self.min_length {
             if str.len() < min_length {
diff --git a/src/validators/typed_dict.rs b/src/validators/typed_dict.rs
@@ -313,7 +313,7 @@ impl Validator for TypedDictValidator {
                             }
                             Err(err) => return Err(err),
                         };
-                        if used_keys.contains(either_str.as_cow().as_ref()) {
+                        if used_keys.contains(either_str.as_cow()?.as_ref()) {
                             continue;
                         }
 
diff --git a/src/validators/union.rs b/src/validators/union.rs
@@ -134,8 +134,8 @@ impl Discriminator {
     fn new(py: Python, raw: &PyAny) -> PyResult<Self> {
         if raw.is_callable() {
             return Ok(Self::Function(raw.to_object(py)));
-        } else if let Ok(str) = raw.strict_str() {
-            if str.as_cow().as_ref() == "self-schema-discriminator" {
+        } else if let Ok(py_str) = raw.cast_as::<PyString>() {
+            if py_str.to_str()? == "self-schema-discriminator" {
                 return Ok(Self::SelfSchema);
             }
         }
@@ -246,7 +246,7 @@ impl Validator for TaggedUnionValidator {
                     GenericMapping::PyGetAttr(obj) => find_validator!(obj, py_get_attr),
                     GenericMapping::JsonObject(mapping) => find_validator!(mapping, json_get),
                 }?;
-                self.find_call_validator(py, tag.as_cow(), input, extra, slots, recursion_guard)
+                self.find_call_validator(py, tag.as_cow()?, input, extra, slots, recursion_guard)
             }
             Discriminator::Function(ref func) => {
                 let tag = func.call1(py, (input.to_object(py),))?;
@@ -301,7 +301,7 @@ impl TaggedUnionValidator {
                 },
                 _ => unreachable!(),
             };
-            let tag_cow = either_tag.as_cow();
+            let tag_cow = either_tag.as_cow()?;
             let tag = tag_cow.as_ref();
             // custom logic to distinguish between different function and tuple schemas
             if tag == "function" || tag == "tuple" {
@@ -314,13 +314,13 @@ impl TaggedUnionValidator {
                 };
                 if tag == "function" {
                     let mode = mode.ok_or_else(|| self.tag_not_found(input))?;
-                    if mode.as_cow().as_ref() == "plain" {
+                    if mode.as_cow()?.as_ref() == "plain" {
                         return Ok(Cow::Borrowed("function-plain"));
                     }
                 } else {
                     // tag == "tuple"
                     if let Some(mode) = mode {
-                        if mode.as_cow().as_ref() == "positional" {
+                        if mode.as_cow()?.as_ref() == "positional" {
                             return Ok(Cow::Borrowed("tuple-positional"));
                         }
                     }
diff --git a/tests/validators/test_bytes.py b/tests/validators/test_bytes.py
@@ -29,6 +29,20 @@ def test_lax_bytes_validator():
 
     assert v.validate_json('"foo"') == b'foo'
 
+    assert v.validate_python('🐈 Hello') == b'\xf0\x9f\x90\x88 Hello'
+    # `.to_str()` Returns a `UnicodeEncodeError` if the input is not valid unicode (containing unpaired surrogates).
+    # https://github.com/PyO3/pyo3/blob/6503128442b8f3e767c663a6a8d96376d7fb603d/src/types/string.rs#L477
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python('🐈 Hello \ud800World')
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'str_unicode',
+            'loc': [],
+            'message': 'Input should be a valid string, unable to parse raw data as a unicode string',
+            'input_value': '🐈 Hello \ud800World',
+        }
+    ]
+
 
 @pytest.mark.parametrize(
     'opts,input,expected',
diff --git a/tests/validators/test_string.py b/tests/validators/test_string.py
@@ -33,6 +33,7 @@ def test_str(py_and_json: PyAndJson, input_value, expected):
     'input_value,expected',
     [
         ('foobar', 'foobar'),
+        ('🐈 Hello \ud800World', '🐈 Hello \ud800World'),
         (b'foobar', 'foobar'),
         (bytearray(b'foobar'), 'foobar'),
         (
@@ -80,6 +81,7 @@ def test_str_not_json(input_value, expected):
         # to_upper and strip comes after pattern check
         ({'to_upper': True, 'pattern': 'abc'}, 'abc', 'ABC'),
         ({'strip_whitespace': True, 'pattern': r'\d+$'}, 'foobar 123 ', Err("String should match pattern '\\d+$'")),
+        ({'min_length': 1}, '🐈 Hello', '🐈 Hello'),
     ],
 )
 def test_constrained_str(py_and_json: PyAndJson, kwargs: Dict[str, Any], input_value, expected):
@@ -91,6 +93,24 @@ def test_constrained_str(py_and_json: PyAndJson, kwargs: Dict[str, Any], input_v
         assert v.validate_test(input_value) == expected
 
 
+def test_unicode_error():
+    # `.to_str()` Returns a `UnicodeEncodeError` if the input is not valid unicode (containing unpaired surrogates).
+    # https://github.com/PyO3/pyo3/blob/6503128442b8f3e767c663a6a8d96376d7fb603d/src/types/string.rs#L477
+    v = SchemaValidator({'type': 'str', 'min_length': 1})
+    assert v.validate_python('🐈 Hello') == '🐈 Hello'
+
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_python('🐈 Hello \ud800World')
+    assert exc_info.value.errors() == [
+        {
+            'kind': 'str_unicode',
+            'loc': [],
+            'message': 'Input should be a valid string, unable to parse raw data as a unicode string',
+            'input_value': '🐈 Hello \ud800World',
+        }
+    ]
+
+
 def test_str_constrained():
     v = SchemaValidator({'type': 'str', 'max_length': 5})
     assert v.validate_python('test') == 'test'

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,7 @@ use std::borrow::Cow;`
`3`	`3`	`use pyo3::prelude::*;`
`4`	`4`	`use pyo3::types::{PyBytes, PyDict, PyFrozenSet, PyList, PySet, PyString, PyTuple};`
`5`	`5`
`6`		`-use crate::errors::{ErrorKind, ValError, ValLineError, ValResult};`
	`6`	`+use crate::errors::{ErrorKind, InputValue, ValError, ValLineError, ValResult};`
`7`	`7`	`use crate::recursion_guard::RecursionGuard;`
`8`	`8`	`use crate::validators::{CombinedValidator, Extra, Validator};`
`9`	`9`
`@@ -212,10 +212,10 @@ pub enum EitherString<'a> {`
`212`	`212`	`}`
`213`	`213`
`214`	`214`	`impl<'a> EitherString<'a> {`
`215`		`- pub fn as_cow(&self) -> Cow<str> {`
	`215`	`+ pub fn as_cow(&self) -> ValResult<'a, Cow<str>> {`
`216`	`216`	`match self {`
`217`		`- Self::Cow(data) => data.clone(),`
`218`		`- Self::Py(py_str) => py_str.to_string_lossy(),`
	`217`	`+ Self::Cow(data) => Ok(data.clone()),`
	`218`	`+ Self::Py(py_str) => Ok(Cow::Borrowed(py_string_str(py_str)?)),`
`219`	`219`	`}`
`220`	`220`	`}`
`221`	`221`
`@@ -251,6 +251,12 @@ impl<'a> IntoPy<PyObject> for EitherString<'a> {`
`251`	`251`	`}`
`252`	`252`	`}`
`253`	`253`
	`254`	`+pub fn py_string_str(py_str: &PyString) -> ValResult<&str> {`
	`255`	`+ py_str.to_str().map_err(\|_\| {`
	`256`	`+ ValError::new_custom_input(ErrorKind::StrUnicode, InputValue::PyObject(py_str.into_py(py_str.py())))`
	`257`	`+ })`
	`258`	`+}`
	`259`	`+`
`254`	`260`	`#[cfg_attr(debug_assertions, derive(Debug))]`
`255`	`261`	`pub enum EitherBytes<'a> {`
`256`	`262`	`Cow(Cow<'a, [u8]>),`
Original file line number	Diff line number	Diff line change
`@@ -281,7 +281,7 @@ impl Validator for ArgumentsValidator {`
`281`	`281`	`}`
`282`	`282`	`Err(err) => return Err(err),`
`283`	`283`	`};`
`284`		`- if !used_kwargs.contains(either_str.as_cow().as_ref()) {`
	`284`	`+ if !used_kwargs.contains(either_str.as_cow()?.as_ref()) {`
`285`	`285`	`match self.var_kwargs_validator {`
`286`	`286`	`Some(ref validator) => match validator.validate(py, value, extra, slots, recursion_guard) {`
`287`	`287`	`Ok(value) => output_kwargs.set_item(either_str.as_py_string(py), value)?,`
Original file line number	Diff line number	Diff line change
`@@ -313,7 +313,7 @@ impl Validator for TypedDictValidator {`
`313`	`313`	`}`
`314`	`314`	`Err(err) => return Err(err),`
`315`	`315`	`};`
`316`		`- if used_keys.contains(either_str.as_cow().as_ref()) {`
	`316`	`+ if used_keys.contains(either_str.as_cow()?.as_ref()) {`
`317`	`317`	`continue;`
`318`	`318`	`}`
`319`	`319`