count recursion count

Author: MarkusSintonen

src/schema_traverse.rs

@@ -1,9 +1,9 @@
+use std::collections::HashMap;
 use crate::tools::py_err;
 use pyo3::exceptions::{PyKeyError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PySet, PyString, PyTuple};
 use pyo3::{intern, Bound, PyResult};
-use std::collections::HashSet;
 
 const CORE_SCHEMA_METADATA_DISCRIMINATOR_PLACEHOLDER_KEY: &str = "pydantic.internal.union_discriminator";
 
@@ -40,13 +40,13 @@ fn gather_definition_ref(schema_ref_dict: &Bound<'_, PyDict>, ctx: &mut GatherCt
         let schema_ref_pystr = schema_ref.downcast_exact::<PyString>()?;
         let schema_ref_str = schema_ref_pystr.to_str()?;
 
-        if !ctx.seen_refs.contains(schema_ref_str) {
+        if *ctx.refs_recursion_count.entry(schema_ref_str.to_string()).or_insert(0) == 0 {
             defaultdict_list_append!(&ctx.def_refs, schema_ref_pystr, schema_ref_dict);
 
             if let Some(def) = ctx.definitions_dict.get_item(schema_ref_pystr)? {
-                ctx.seen_refs.insert(schema_ref_str.to_string());
+                *ctx.refs_recursion_count.get_mut(schema_ref_str).unwrap() += 1;
                 gather_schema(def.downcast_exact::<PyDict>()?, ctx)?;
-                ctx.seen_refs.remove(schema_ref_str);
+                *ctx.refs_recursion_count.get_mut(schema_ref_str).unwrap() -= 1;
             }
             Ok(false)
         } else {
@@ -183,7 +183,7 @@ pub struct GatherCtx<'a, 'py> {
     pub def_refs: Bound<'py, PyDict>,
     pub recursive_def_refs: Bound<'py, PySet>,
     pub discriminators: Bound<'py, PyList>,
-    seen_refs: HashSet<String>,
+    refs_recursion_count: HashMap<String, i32>,
 }
 
 impl<'a, 'py> GatherCtx<'a, 'py> {
@@ -193,7 +193,7 @@ impl<'a, 'py> GatherCtx<'a, 'py> {
             def_refs: PyDict::new_bound(definitions.py()),
             recursive_def_refs: PySet::empty_bound(definitions.py())?,
             discriminators: PyList::empty_bound(definitions.py()),
-            seen_refs: HashSet::new(),
+            refs_recursion_count: HashMap::default(),
         };
         Ok(ctx)
     }