11use std:: borrow:: Cow ;
22use std:: fmt:: Debug ;
3+ use std:: io:: { self , Write } ;
34
45use pyo3:: exceptions:: PyTypeError ;
56use pyo3:: prelude:: * ;
@@ -9,7 +10,7 @@ use pyo3::{intern, PyTraverseError, PyVisit};
910
1011use enum_dispatch:: enum_dispatch;
1112use serde:: Serialize ;
12- use serde_json:: ser:: PrettyFormatter ;
13+ use serde_json:: ser:: { Formatter , PrettyFormatter } ;
1314
1415use crate :: build_tools:: py_schema_err;
1516use crate :: build_tools:: py_schema_error_type;
@@ -432,6 +433,71 @@ impl Serialize for PydanticSerializer<'_> {
432433 }
433434}
434435
436+ struct EscapeNonAsciiFormatter ;
437+
438+ impl Formatter for EscapeNonAsciiFormatter {
439+ fn write_string_fragment < W : ?Sized + Write > ( & mut self , writer : & mut W , fragment : & str ) -> io:: Result < ( ) > {
440+ for ch in fragment. chars ( ) {
441+ if ch. is_ascii ( ) {
442+ writer. write_all ( ch. encode_utf8 ( & mut [ 0 ; 4 ] ) . as_bytes ( ) ) ?;
443+ } else {
444+ for escape in ch. encode_utf16 ( & mut [ 0 ; 2 ] ) {
445+ write ! ( writer, "\\ u{escape:04x}" ) ?;
446+ }
447+ }
448+ }
449+ Ok ( ( ) )
450+ }
451+ }
452+
453+ struct EscapeNonAsciiPrettyFormatter < ' a > {
454+ pretty : PrettyFormatter < ' a > ,
455+ escape_non_ascii : EscapeNonAsciiFormatter ,
456+ }
457+
458+ impl < ' a > EscapeNonAsciiPrettyFormatter < ' a > {
459+ pub fn with_indent ( indent : & ' a [ u8 ] ) -> Self {
460+ Self {
461+ pretty : PrettyFormatter :: with_indent ( indent) ,
462+ escape_non_ascii : EscapeNonAsciiFormatter ,
463+ }
464+ }
465+ }
466+
467+ macro_rules! defer {
468+ ( $formatter: ident, $fun: ident) => {
469+ fn $fun<W >( & mut self , writer: & mut W ) -> io:: Result <( ) >
470+ where
471+ W : ?Sized + io:: Write ,
472+ {
473+ self . $formatter. $fun( writer)
474+ }
475+ } ;
476+ ( $formatter: ident, $fun: ident, $val: ty) => {
477+ fn $fun<W >( & mut self , writer: & mut W , val: $val) -> io:: Result <( ) >
478+ where
479+ W : ?Sized + io:: Write ,
480+ {
481+ self . $formatter. $fun( writer, val)
482+ }
483+ } ;
484+ }
485+
486+ #[ allow( clippy:: needless_lifetimes) ]
487+ impl < ' a > Formatter for EscapeNonAsciiPrettyFormatter < ' a > {
488+ defer ! ( escape_non_ascii, write_string_fragment, & str ) ;
489+ defer ! ( pretty, begin_array) ;
490+ defer ! ( pretty, end_array) ;
491+ defer ! ( pretty, begin_array_value, bool ) ;
492+ defer ! ( pretty, end_array_value) ;
493+ defer ! ( pretty, begin_object) ;
494+ defer ! ( pretty, end_object) ;
495+ defer ! ( pretty, begin_object_key, bool ) ;
496+ defer ! ( pretty, end_object_key) ;
497+ defer ! ( pretty, begin_object_value) ;
498+ defer ! ( pretty, end_object_value) ;
499+ }
500+
435501#[ allow( clippy:: too_many_arguments) ]
436502pub ( crate ) fn to_json_bytes (
437503 value : & Bound < ' _ , PyAny > ,
@@ -440,25 +506,40 @@ pub(crate) fn to_json_bytes(
440506 exclude : Option < & Bound < ' _ , PyAny > > ,
441507 extra : & Extra ,
442508 indent : Option < usize > ,
509+ ensure_ascii : bool ,
443510 expected_json_size : usize ,
444511) -> PyResult < Vec < u8 > > {
445512 let serializer = PydanticSerializer :: new ( value, serializer, include, exclude, extra) ;
446513
447514 let writer: Vec < u8 > = Vec :: with_capacity ( expected_json_size) ;
448- let bytes = match indent {
449- Some ( indent) => {
515+
516+ let bytes = match ( indent, ensure_ascii) {
517+ ( Some ( indent) , true ) => {
518+ let indent = vec ! [ b' ' ; indent] ;
519+ let formatter = EscapeNonAsciiPrettyFormatter :: with_indent ( & indent) ;
520+ let mut ser = PythonSerializer :: with_formatter ( writer, formatter) ;
521+ serializer. serialize ( & mut ser) . map_err ( se_err_py_err) ?;
522+ ser. into_inner ( )
523+ }
524+ ( Some ( indent) , false ) => {
450525 let indent = vec ! [ b' ' ; indent] ;
451526 let formatter = PrettyFormatter :: with_indent ( & indent) ;
452527 let mut ser = PythonSerializer :: with_formatter ( writer, formatter) ;
453528 serializer. serialize ( & mut ser) . map_err ( se_err_py_err) ?;
454529 ser. into_inner ( )
455530 }
456- None => {
531+ ( None , true ) => {
532+ let mut ser = PythonSerializer :: with_formatter ( writer, EscapeNonAsciiFormatter ) ;
533+ serializer. serialize ( & mut ser) . map_err ( se_err_py_err) ?;
534+ ser. into_inner ( )
535+ }
536+ ( None , false ) => {
457537 let mut ser = PythonSerializer :: new ( writer) ;
458538 serializer. serialize ( & mut ser) . map_err ( se_err_py_err) ?;
459539 ser. into_inner ( )
460540 }
461541 } ;
542+
462543 Ok ( bytes)
463544}
464545
0 commit comments