Add support for type in json schema to be a list (#138)

Close #126

---------

Signed-off-by: Benjamin <[email protected]>

Author: bnjjj

src/error.rs

@@ -57,8 +57,8 @@ pub enum Error {
     ExternalReferencesNotSupported(Box<str>),
     #[error("Invalid reference format: {0}")]
     InvalidReferenceFormat(Box<str>),
-    #[error("'type' must be a string")]
-    TypeMustBeAString,
+    #[error("'type' must be a string or an array of string")]
+    TypeMustBeAStringOrArray,
     #[error("Unsupported type: {0}")]
     UnsupportedType(Box<str>),
     #[error("maxLength must be greater than or equal to minLength")]
@@ -80,7 +80,8 @@ impl Error {
 #[cfg(feature = "python-bindings")]
 impl From<Error> for pyo3::PyErr {
     fn from(e: Error) -> Self {
-        use pyo3::{exceptions::PyValueError, PyErr};
+        use pyo3::exceptions::PyValueError;
+        use pyo3::PyErr;
         PyErr::new::<PyValueError, _>(e.to_string())
     }
 }

src/json_schema/mod.rs

@@ -100,6 +100,7 @@
 //! An empty object means unconstrained, allowing any JSON type.
 
 use serde_json::Value;
+pub use types::*;
 
 mod parsing;
 pub mod types;
@@ -189,9 +190,10 @@ pub fn regex_from_value(json: &Value, whitespace_pattern: Option<&str>) -> Resul
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use regex::Regex;
 
+    use super::*;
+
     fn should_match(re: &Regex, value: &str) {
         // Asserts that value is fully matched.
         match re.find(value) {
@@ -1084,6 +1086,25 @@ mod tests {
                     "[email protected]",         // multiple errors in domain
                 ]
             ),
+
+            // ==========================================================
+            //                      Multiple types
+            // ==========================================================
+            (
+                r#"{
+                    "title": "Foo",
+                    "type": ["string", "number", "boolean"]
+                }"#,
+                format!(r#"((?:"{STRING_INNER}*")|(?:{NUMBER})|(?:{BOOLEAN}))"#).as_str(),
+                vec!["12.3", "true", r#""a""#],
+                vec![
+                    "null",
+                    "",
+                    "12true",
+                    r#"1.3"a""#,
+                    r#"12.3true"a""#,
+                ],
+            ),
             // Confirm that oneOf doesn't produce illegal lookaround: https://github.com/dottxt-ai/outlines/issues/823
             //
             // The pet field uses the discriminator field to decide which schema (Cat or Dog) applies, based on the pet_type property.

src/json_schema/parsing.rs

@@ -7,7 +7,8 @@ use serde_json::json;
 use serde_json::Value;
 
 use crate::json_schema::types;
-use crate::{Error, Result};
+use crate::Error;
+use crate::Result;
 
 pub(crate) struct Parser<'a> {
     root: &'a Value,
@@ -305,9 +306,39 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_type(&mut self, obj: &serde_json::Map<String, Value>) -> Result<String> {
-        let instance_type = obj["type"]
-            .as_str()
-            .ok_or_else(|| Error::TypeMustBeAString)?;
+        match obj.get("type") {
+            Some(Value::String(instance_type)) => self.parse_type_string(instance_type, obj),
+            Some(Value::Array(instance_types)) => self.parse_type_array(instance_types, obj),
+            _ => Err(Error::TypeMustBeAStringOrArray),
+        }
+    }
+
+    fn parse_type_array(
+        &mut self,
+        instance_types: &[serde_json::Value],
+        obj: &serde_json::Map<String, Value>,
+    ) -> Result<String> {
+        let xor_patterns = instance_types
+            .iter()
+            .map(|instance_type| match instance_type.as_str() {
+                Some(instance_type) => {
+                    let sub_regex = self.parse_type_string(instance_type, obj)?;
+
+                    Ok(format!(r"(?:{})", sub_regex))
+                }
+                None => Err(Error::TypeMustBeAStringOrArray),
+            })
+            .collect::<Result<Vec<String>>>()?
+            .join("|");
+
+        Ok(format!(r"({})", xor_patterns))
+    }
+
+    fn parse_type_string(
+        &mut self,
+        instance_type: &str,
+        obj: &serde_json::Map<String, Value>,
+    ) -> Result<String> {
         match instance_type {
             "string" => self.parse_string_type(obj),
             "number" => self.parse_number_type(obj),