11// Validator for things inside of a typing.Literal[]
22// which can be an int, a string, bytes or an Enum value (including `class Foo(str, Enum)` type enums)
3+ use core:: fmt:: Debug ;
34
4- use ahash:: AHashSet ;
5+ use ahash:: AHashMap ;
56use pyo3:: intern;
67use pyo3:: prelude:: * ;
78use pyo3:: types:: { PyDict , PyList } ;
@@ -15,15 +16,96 @@ use crate::tools::SchemaDict;
1516use super :: { BuildValidator , CombinedValidator , Definitions , DefinitionsBuilder , Extra , Validator } ;
1617
1718#[ derive( Debug , Clone ) ]
18- pub struct LiteralValidator {
19+ pub struct LiteralLookup < T : Clone + Debug > {
1920 // Specialized lookups for ints and strings because they
2021 // (1) are easy to convert between Rust and Python
2122 // (2) hashing them in Rust is very fast
2223 // (3) are the most commonly used things in Literal[...]
23- expected_int : Option < AHashSet < i64 > > ,
24- expected_str : Option < AHashSet < String > > ,
24+ expected_int : Option < AHashMap < i64 , usize > > ,
25+ expected_str : Option < AHashMap < String , usize > > ,
2526 // Catch all for Enum and bytes (the latter only because it is seldom used)
2627 expected_py : Option < Py < PyDict > > ,
28+ pub values : Vec < T > ,
29+ }
30+
31+ impl < T : Clone + Debug > LiteralLookup < T > {
32+ pub fn new < ' py > ( py : Python < ' py > , expected : impl Iterator < Item = ( & ' py PyAny , T ) > ) -> PyResult < Self > {
33+ let mut expected_int = AHashMap :: new ( ) ;
34+ let mut expected_str = AHashMap :: new ( ) ;
35+ let expected_py = PyDict :: new ( py) ;
36+ let mut values = Vec :: new ( ) ;
37+ for ( k, v) in expected {
38+ let id = values. len ( ) ;
39+ values. push ( v) ;
40+ if let Ok ( either_int) = k. exact_int ( ) {
41+ let int = either_int
42+ . into_i64 ( py)
43+ . map_err ( |_| py_schema_error_type ! ( "error extracting int {:?}" , k) ) ?;
44+ expected_int. insert ( int, id) ;
45+ } else if let Ok ( either_str) = k. exact_str ( ) {
46+ let str = either_str
47+ . as_cow ( )
48+ . map_err ( |_| py_schema_error_type ! ( "error extracting str {:?}" , k) ) ?;
49+ expected_str. insert ( str. to_string ( ) , id) ;
50+ } else {
51+ expected_py. set_item ( k, id) ?;
52+ }
53+ }
54+
55+ Ok ( Self {
56+ expected_int : match expected_int. is_empty ( ) {
57+ true => None ,
58+ false => Some ( expected_int) ,
59+ } ,
60+ expected_str : match expected_str. is_empty ( ) {
61+ true => None ,
62+ false => Some ( expected_str) ,
63+ } ,
64+ expected_py : match expected_py. is_empty ( ) {
65+ true => None ,
66+ false => Some ( expected_py. into ( ) ) ,
67+ } ,
68+ values,
69+ } )
70+ }
71+
72+ pub fn validate < ' data , I : Input < ' data > > (
73+ & self ,
74+ py : Python < ' data > ,
75+ input : & ' data I ,
76+ ) -> ValResult < ' data , Option < ( & ' data I , & T ) > > {
77+ // dbg!(input.to_object(py).as_ref(py).repr().unwrap());
78+ if let Some ( expected_ints) = & self . expected_int {
79+ if let Ok ( either_int) = input. exact_int ( ) {
80+ let int = either_int. into_i64 ( py) ?;
81+ if let Some ( id) = expected_ints. get ( & int) {
82+ return Ok ( Some ( ( input, & self . values [ * id] ) ) ) ;
83+ }
84+ }
85+ }
86+ if let Some ( expected_strings) = & self . expected_str {
87+ // dbg!(expected_strings);
88+ if let Ok ( either_str) = input. exact_str ( ) {
89+ let cow = either_str. as_cow ( ) ?;
90+ if let Some ( id) = expected_strings. get ( cow. as_ref ( ) ) {
91+ return Ok ( Some ( ( input, & self . values [ * id] ) ) ) ;
92+ }
93+ }
94+ }
95+ // must be an enum or bytes
96+ if let Some ( expected_py) = & self . expected_py {
97+ if let Some ( v) = expected_py. as_ref ( py) . get_item ( input) {
98+ let id: usize = v. extract ( ) . unwrap ( ) ;
99+ return Ok ( Some ( ( input, & self . values [ id] ) ) ) ;
100+ }
101+ } ;
102+ Ok ( None )
103+ }
104+ }
105+
106+ #[ derive( Debug , Clone ) ]
107+ pub struct LiteralValidator {
108+ lookup : LiteralLookup < PyObject > ,
27109 expected_repr : String ,
28110 name : String ,
29111}
@@ -41,32 +123,14 @@ impl BuildValidator for LiteralValidator {
41123 return py_schema_err ! ( "`expected` should have length > 0" ) ;
42124 }
43125 let py = expected. py ( ) ;
44- // Literal[...] only supports int, str, bytes or enums, all of which can be hashed
45- let mut expected_int = AHashSet :: new ( ) ;
46- let mut expected_str = AHashSet :: new ( ) ;
47- let expected_py = PyDict :: new ( py) ;
48126 let mut repr_args: Vec < String > = Vec :: new ( ) ;
49127 for item in expected. iter ( ) {
50128 repr_args. push ( item. repr ( ) ?. extract ( ) ?) ;
51- if let Ok ( either_int) = item. strict_int ( ) {
52- let int = either_int
53- . into_i64 ( py)
54- . map_err ( |_| py_schema_error_type ! ( "error extracting int {:?}" , item) ) ?;
55- expected_int. insert ( int) ;
56- } else if let Ok ( either_str) = item. strict_str ( ) {
57- let str = either_str
58- . as_cow ( )
59- . map_err ( |_| py_schema_error_type ! ( "error extracting str {:?}" , item) ) ?;
60- expected_str. insert ( str. to_string ( ) ) ;
61- } else {
62- expected_py. set_item ( item, item) ?;
63- }
64129 }
65130 let ( expected_repr, name) = expected_repr_name ( repr_args, "literal" ) ;
131+ let lookup = LiteralLookup :: new ( py, expected. iter ( ) . map ( |v| ( v, v. to_object ( py) ) ) ) ?;
66132 Ok ( CombinedValidator :: Literal ( Self {
67- expected_int : ( !expected_int. is_empty ( ) ) . then_some ( expected_int) ,
68- expected_str : ( !expected_str. is_empty ( ) ) . then_some ( expected_str) ,
69- expected_py : ( !expected_py. is_empty ( ) ) . then_some ( expected_py. into ( ) ) ,
133+ lookup,
70134 expected_repr,
71135 name,
72136 } ) )
@@ -82,34 +146,15 @@ impl Validator for LiteralValidator {
82146 _definitions : & ' data Definitions < CombinedValidator > ,
83147 _recursion_guard : & ' s mut RecursionGuard ,
84148 ) -> ValResult < ' data , PyObject > {
85- if let Some ( expected_ints) = & self . expected_int {
86- if let Ok ( either_int) = input. strict_int ( ) {
87- let int = either_int. into_i64 ( py) ?;
88- if expected_ints. contains ( & int) {
89- return Ok ( input. to_object ( py) ) ;
90- }
91- }
149+ match self . lookup . validate ( py, input) ? {
150+ Some ( ( _, v) ) => Ok ( v. clone ( ) ) ,
151+ None => Err ( ValError :: new (
152+ ErrorType :: LiteralError {
153+ expected : self . expected_repr . clone ( ) ,
154+ } ,
155+ input,
156+ ) ) ,
92157 }
93- if let Some ( expected_strings) = & self . expected_str {
94- if let Ok ( either_str) = input. strict_str ( ) {
95- let cow = either_str. as_cow ( ) ?;
96- if expected_strings. contains ( cow. as_ref ( ) ) {
97- return Ok ( input. to_object ( py) ) ;
98- }
99- }
100- }
101- // must be an enum or bytes
102- if let Some ( expected_py) = & self . expected_py {
103- if let Some ( v) = expected_py. as_ref ( py) . get_item ( input) {
104- return Ok ( v. into ( ) ) ;
105- }
106- } ;
107- Err ( ValError :: new (
108- ErrorType :: LiteralError {
109- expected : self . expected_repr . clone ( ) ,
110- } ,
111- input,
112- ) )
113158 }
114159
115160 fn different_strict_behavior (
0 commit comments