fix: Make JSON Schema checks actually work again

hugr-core/src/hugr/serialize/test.rs

@@ -62,34 +62,37 @@ impl NamedSchema {
         Self { name, schema }
     }
 
-    pub fn check(&self, val: &serde_json::Value) {
+    pub fn check(&self, val: &serde_json::Value) -> Result<(), String> {
         let mut errors = self.schema.iter_errors(val).peekable();
-        if errors.peek().is_some() {
-            // errors don't necessarily implement Debug
-            eprintln!("Schema failed to validate: {}", self.name);
-            for error in errors {
-                eprintln!("Validation error: {error}");
-                eprintln!("Instance path: {}", error.instance_path);
-            }
-            panic!("Serialization test failed.");
+        if errors.peek().is_none() {
+            return Ok(());
         }
+
+        // errors don't necessarily implement Debug
+        let mut strs = vec![format!("Schema failed to validate: {}", self.name)];
+        strs.extend(errors.flat_map(|error| {
+            [
+                format!("Validation error: {error}"),
+                format!("Instance path: {}", error.instance_path),
+            ]
+        }));
+        strs.push("Serialization test failed.".to_string());
+        Err(strs.join("\n"))
     }
 
     pub fn check_schemas(
         val: &serde_json::Value,
         schemas: impl IntoIterator<Item = &'static Self>,
-    ) {
-        for schema in schemas {
-            schema.check(val);
-        }
+    ) -> Result<(), String> {
+        schemas.into_iter().try_for_each(|schema| schema.check(val))
     }
 }
 
 macro_rules! include_schema {
     ($name:ident, $path:literal) => {
         lazy_static! {
             static ref $name: NamedSchema =
-                NamedSchema::new("$name", {
+                NamedSchema::new(stringify!($name), {
                     let schema_val: serde_json::Value = serde_json::from_str(include_str!(
                         concat!("../../../../specification/schema/", $path, "_live.json")
                     ))
@@ -161,7 +164,7 @@ fn ser_deserialize_check_schema<T: serde::de::DeserializeOwned>(
     val: serde_json::Value,
     schemas: impl IntoIterator<Item = &'static NamedSchema>,
 ) -> T {
-    NamedSchema::check_schemas(&val, schemas);
+    NamedSchema::check_schemas(&val, schemas).unwrap();
     serde_json::from_value(val).unwrap()
 }
 
@@ -171,8 +174,10 @@ fn ser_roundtrip_check_schema<TSer: Serialize, TDeser: serde::de::DeserializeOwn
     schemas: impl IntoIterator<Item = &'static NamedSchema>,
 ) -> TDeser {
     let val = serde_json::to_value(g).unwrap();
-    NamedSchema::check_schemas(&val, schemas);
-    serde_json::from_value(val).unwrap()
+    match NamedSchema::check_schemas(&val, schemas) {
+        Ok(()) => serde_json::from_value(val).unwrap(),
+        Err(msg) => panic!("ser_roundtrip_check_schema failed with {msg}, input was {val}"),
+    }
 }
 
 /// Serialize a Hugr and check that it is valid against the schema.
@@ -187,7 +192,7 @@ pub(crate) fn check_hugr_serialization_schema(hugr: &Hugr) {
     let schemas = get_schemas(true);
     let hugr_ser = HugrSer(hugr);
     let val = serde_json::to_value(hugr_ser).unwrap();
-    NamedSchema::check_schemas(&val, schemas);
+    NamedSchema::check_schemas(&val, schemas).unwrap();
 }
 
 /// Serialize and deserialize a HUGR, and check that the result is the same as the original.
@@ -225,6 +230,77 @@ fn check_testing_roundtrip(t: impl Into<SerTestingLatest>) {
     assert_eq!(before, after);
 }
 
+fn test_schema_val() -> serde_json::Value {
+    serde_json::json!({
+        "op_def":null,
+        "optype":{
+            "name":"polyfunc1",
+            "op":"FuncDefn",
+            "parent":0,
+            "signature":{
+                "body":{
+                    "input":[],
+                    "output":[]
+                },
+                "params":[
+                    {"bound":null,"tp":"BoundedNat"}
+                ]
+            }
+        },
+        "poly_func_type":null,
+        "sum_type":null,
+        "typ":null,
+        "value":null,
+        "version":"live"
+    })
+}
+
+fn schema_val() -> serde_json::Value {
+    serde_json::json!({"nodes": [], "edges": [], "version": "live"})
+}
+
+#[rstest]
+#[case(&TESTING_SCHEMA, &TESTING_SCHEMA_STRICT, test_schema_val(), Some("optype"))]
+#[case(&SCHEMA, &SCHEMA_STRICT, schema_val(), None)]
+fn wrong_fields(
+    #[case] lax_schema: &'static NamedSchema,
+    #[case] strict_schema: &'static NamedSchema,
+    #[case] mut val: serde_json::Value,
+    #[case] target_loc: impl IntoIterator<Item = &'static str> + Clone,
+) {
+    use serde_json::Value;
+    fn get_fields(
+        val: &mut Value,
+        mut path: impl Iterator<Item = &'static str>,
+    ) -> &mut serde_json::Map<String, Value> {
+        let Value::Object(fields) = val else { panic!() };
+        match path.next() {
+            Some(n) => get_fields(fields.get_mut(n).unwrap(), path),
+            None => fields,
+        }
+    }
+    // First, some "known good" JSON
+    NamedSchema::check_schemas(&val, [lax_schema, strict_schema]).unwrap();
+
+    // Now try adding an extra field
+    let fields = get_fields(&mut val, target_loc.clone().into_iter());
+    fields.insert(
+        "extra_field".to_string(),
+        Value::String("not in schema".to_string()),
+    );
+    strict_schema.check(&val).unwrap_err();
+    lax_schema.check(&val).unwrap();
+
+    // And removing one
+    let fields = get_fields(&mut val, target_loc.into_iter());
+    fields.remove("extra_field").unwrap();
+    let key = fields.keys().next().unwrap().clone();
+    fields.remove(&key).unwrap();
+
+    lax_schema.check(&val).unwrap_err();
+    strict_schema.check(&val).unwrap_err();
+}
+
 /// Generate an optype for a node with a matching amount of inputs and outputs.
 fn gen_optype(g: &MultiPortGraph, node: portgraph::NodeIndex) -> OpType {
     let inputs = g.num_inputs(node);
@@ -544,7 +620,7 @@ fn std_extensions_valid() {
     let std_reg = crate::std_extensions::std_reg();
     for ext in std_reg {
         let val = serde_json::to_value(ext).unwrap();
-        NamedSchema::check_schemas(&val, get_schemas(true));
+        NamedSchema::check_schemas(&val, get_schemas(true)).unwrap();
         // check deserialises correctly, can't check equality because of custom binaries.
         let deser: crate::extension::Extension = serde_json::from_value(val.clone()).unwrap();
         assert_eq!(serde_json::to_value(deser).unwrap(), val);

hugr-core/src/ops/custom.rs

@@ -7,7 +7,8 @@ use thiserror::Error;
 #[cfg(test)]
 use {
     crate::extension::test::SimpleOpDef, crate::proptest::any_nonempty_smolstr,
-    ::proptest::prelude::*, ::proptest_derive::Arbitrary,
+    crate::types::proptest_utils::any_serde_type_arg_vec, ::proptest::prelude::*,
+    ::proptest_derive::Arbitrary,
 };
 
 use crate::core::HugrNode;
@@ -35,6 +36,7 @@ pub struct ExtensionOp {
         proptest(strategy = "any::<SimpleOpDef>().prop_map(|x| Arc::new(x.into()))")
     )]
     def: Arc<OpDef>,
+    #[cfg_attr(test, proptest(strategy = "any_serde_type_arg_vec()"))]
     args: Vec<TypeArg>,
     signature: Signature, // Cache
 }
@@ -235,6 +237,7 @@ pub struct OpaqueOp {
     extension: ExtensionId,
     #[cfg_attr(test, proptest(strategy = "any_nonempty_smolstr()"))]
     name: OpName,
+    #[cfg_attr(test, proptest(strategy = "any_serde_type_arg_vec()"))]
     args: Vec<TypeArg>,
     // note that the `signature` field might not include `extension`. Thus this must
     // remain private, and should be accessed through

hugr-core/src/ops/dataflow.rs

@@ -10,7 +10,7 @@ use crate::types::{EdgeKind, PolyFuncType, Signature, Substitution, Type, TypeAr
 use crate::{IncomingPort, type_row};
 
 #[cfg(test)]
-use proptest_derive::Arbitrary;
+use {crate::types::proptest_utils::any_serde_type_arg_vec, proptest_derive::Arbitrary};
 
 /// Trait implemented by all dataflow operations.
 pub trait DataflowOpTrait: Sized {
@@ -191,6 +191,7 @@ pub struct Call {
     /// Signature of function being called.
     pub func_sig: PolyFuncType,
     /// The type arguments that instantiate `func_sig`.
+    #[cfg_attr(test, proptest(strategy = "any_serde_type_arg_vec()"))]
     pub type_args: Vec<TypeArg>,
     /// The instantiation of `func_sig`.
     pub instantiation: Signature, // Cache, so we can fail in try_new() not in signature()
@@ -391,6 +392,7 @@ pub struct LoadFunction {
     /// Signature of the function
     pub func_sig: PolyFuncType,
     /// The type arguments that instantiate `func_sig`.
+    #[cfg_attr(test, proptest(strategy = "any_serde_type_arg_vec()"))]
     pub type_args: Vec<TypeArg>,
     /// The instantiation of `func_sig`.
     pub instantiation: Signature, // Cache, so we can fail in try_new() not in signature()

hugr-core/src/types.rs

@@ -1107,3 +1107,82 @@ pub(crate) mod test {
         }
     }
 }
+
+#[cfg(test)]
+pub(super) mod proptest_utils {
+    use proptest::collection::vec;
+    use proptest::prelude::{Strategy, any_with};
+
+    use super::serialize::{TermSer, TypeArgSer, TypeParamSer};
+    use super::type_param::Term;
+
+    use crate::proptest::RecursionDepth;
+    use crate::types::serialize::ArrayOrTermSer;
+
+    fn term_is_serde_type_arg(t: &Term) -> bool {
+        let TermSer::TypeArg(arg) = TermSer::from(t.clone()) else {
+            return false;
+        };
+        match arg {
+            TypeArgSer::List { elems: terms }
+            | TypeArgSer::ListConcat { lists: terms }
+            | TypeArgSer::Tuple { elems: terms }
+            | TypeArgSer::TupleConcat { tuples: terms } => terms.iter().all(term_is_serde_type_arg),
+            TypeArgSer::Variable { v } => term_is_serde_type_param(&v.cached_decl),
+            TypeArgSer::Type { ty } => {
+                if let Some(cty) = ty.as_extension() {
+                    cty.args().iter().all(term_is_serde_type_arg)
+                } else {
+                    true
+                }
+            } // Do we need to inspect inside function types? sum types?
+            TypeArgSer::BoundedNat { .. }
+            | TypeArgSer::String { .. }
+            | TypeArgSer::Bytes { .. }
+            | TypeArgSer::Float { .. } => true,
+        }
+    }
+
+    fn term_is_serde_type_param(t: &Term) -> bool {
+        let TermSer::TypeParam(parm) = TermSer::from(t.clone()) else {
+            return false;
+        };
+        match parm {
+            TypeParamSer::Type { .. }
+            | TypeParamSer::BoundedNat { .. }
+            | TypeParamSer::String
+            | TypeParamSer::Bytes
+            | TypeParamSer::Float
+            | TypeParamSer::StaticType => true,
+            TypeParamSer::List { param } => term_is_serde_type_param(&param),
+            TypeParamSer::Tuple { params } => {
+                match &params {
+                    ArrayOrTermSer::Array(terms) => terms.iter().all(term_is_serde_type_param),
+                    ArrayOrTermSer::Term(b) => match &**b {
+                        Term::List(_) => panic!("Should be represented as ArrayOrTermSer::Array"),
+                        // This might be well-typed, but does not fit the (TODO: update) JSON schema
+                        Term::Variable(_) => false,
+                        // Similarly, but not produced by our `impl Arbitrary`:
+                        Term::ListConcat(_) => todo!("Update schema"),
+
+                        // The others do not fit the JSON schema, and are not well-typed,
+                        // but can be produced by our impl of Arbitrary, so we must filter out:
+                        _ => false,
+                    },
+                }
+            }
+        }
+    }
+
+    pub fn any_serde_type_arg(depth: RecursionDepth) -> impl Strategy<Value = Term> {
+        any_with::<Term>(depth).prop_filter("Term was not a TypeArg", term_is_serde_type_arg)
+    }
+
+    pub fn any_serde_type_arg_vec() -> impl Strategy<Value = Vec<Term>> {
+        vec(any_serde_type_arg(RecursionDepth::default()), 1..3)
+    }
+
+    pub fn any_serde_type_param(depth: RecursionDepth) -> impl Strategy<Value = Term> {
+        any_with::<Term>(depth).prop_filter("Term was not a TypeParam", term_is_serde_type_param)
+    }
+}

hugr-core/src/types/custom.rs

@@ -188,7 +188,7 @@ mod test {
         use crate::extension::ExtensionId;
         use crate::proptest::RecursionDepth;
         use crate::proptest::any_nonempty_string;
-        use crate::types::type_param::TypeArg;
+        use crate::types::proptest_utils::any_serde_type_arg;
         use crate::types::{CustomType, TypeBound};
         use ::proptest::collection::vec;
         use ::proptest::prelude::*;
@@ -224,7 +224,7 @@ mod test {
                     Just(vec![]).boxed()
                 } else {
                     // a TypeArg may contain a CustomType, so we descend here
-                    vec(any_with::<TypeArg>(depth.descend()), 0..3).boxed()
+                    vec(any_serde_type_arg(depth.descend()), 0..3).boxed()
                 };
                 (any_nonempty_string(), args, any::<ExtensionId>(), bound)
                     .prop_map(|(id, args, extension, bound)| {

hugr-core/src/types/poly_func.rs

@@ -7,6 +7,7 @@ use itertools::Itertools;
 use crate::extension::SignatureError;
 #[cfg(test)]
 use {
+    super::proptest_utils::any_serde_type_param,
     crate::proptest::RecursionDepth,
     ::proptest::{collection::vec, prelude::*},
     proptest_derive::Arbitrary,
@@ -31,7 +32,7 @@ pub struct PolyFuncTypeBase<RV: MaybeRV> {
     /// The declared type parameters, i.e., these must be instantiated with
     /// the same number of [`TypeArg`]s before the function can be called. This
     /// defines the indices used by variables inside the body.
-    #[cfg_attr(test, proptest(strategy = "vec(any_with::<TypeParam>(params), 0..3)"))]
+    #[cfg_attr(test, proptest(strategy = "vec(any_serde_type_param(params), 0..3)"))]
     params: Vec<TypeParam>,
     /// Template for the function. May contain variables up to length of [`Self::params`]
     #[cfg_attr(test, proptest(strategy = "any_with::<FuncTypeBase<RV>>(params)"))]

hugr-core/src/types/serialize.rs

@@ -187,7 +187,7 @@ impl From<TermSer> for Term {
 #[serde(untagged)]
 pub(super) enum ArrayOrTermSer {
     Array(Vec<Term>),
-    Term(Box<Term>),
+    Term(Box<Term>), // TODO JSON Schema does not really support this yet
 }
 
 impl From<ArrayOrTermSer> for Term {

hugr-core/src/types/type_param.rs

@@ -261,7 +261,7 @@ impl<const N: usize> From<[Term; N]> for Term {
 #[display("#{idx}")]
 pub struct TermVar {
     idx: usize,
-    cached_decl: Box<Term>,
+    pub(in crate::types) cached_decl: Box<Term>,
 }
 
 impl Term {
@@ -1046,13 +1046,13 @@ mod test {
 
         use super::super::{TermVar, UpperBound};
         use crate::proptest::RecursionDepth;
-        use crate::types::{Term, Type, TypeBound};
+        use crate::types::{Term, Type, TypeBound, proptest_utils::any_serde_type_param};
 
         impl Arbitrary for TermVar {
             type Parameters = RecursionDepth;
             type Strategy = BoxedStrategy<Self>;
             fn arbitrary_with(depth: Self::Parameters) -> Self::Strategy {
-                (any::<usize>(), any_with::<Term>(depth))
+                (any::<usize>(), any_serde_type_param(depth))
                     .prop_map(|(idx, cached_decl)| Self {
                         idx,
                         cached_decl: Box::new(cached_decl),

scripts/generate_schema.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 
 from pydantic import ConfigDict
-from pydantic.json_schema import models_json_schema
+from pydantic.json_schema import DEFAULT_REF_TEMPLATE, models_json_schema
 
 from hugr._serialization.extension import Extension, Package
 from hugr._serialization.serial_hugr import SerialHugr
@@ -38,6 +38,9 @@ def write_schema(
     _, top_level_schema = models_json_schema(
         [(s, "validation") for s in schemas], title="HUGR schema"
     )
+    top_level_schema["oneOf"] = [
+        {"$ref": DEFAULT_REF_TEMPLATE.format(model=s.__name__)} for s in schemas
+    ]
     with path.open("w") as f:
         json.dump(top_level_schema, f, indent=4)