From a58f697363d0b84c7253cbd5f4b9b45826ad21a4 Mon Sep 17 00:00:00 2001
From: Piotr Siuszko <siuszko@zoho.com>
Date: Sat, 5 Jul 2025 11:43:27 +0200
Subject: [PATCH] Add helper trait for registering schema types

---
 crates/bevy_remote/src/builtin_methods.rs |  7 ++-
 crates/bevy_remote/src/lib.rs             |  6 +--
 crates/bevy_remote/src/schemas/mod.rs     | 57 ++++++++++++++++++++++-
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/crates/bevy_remote/src/builtin_methods.rs b/crates/bevy_remote/src/builtin_methods.rs
index e448355319..2f699bdd2e 100644
--- a/crates/bevy_remote/src/builtin_methods.rs
+++ b/crates/bevy_remote/src/builtin_methods.rs
@@ -1751,7 +1751,12 @@ mod tests {
         world.insert_resource(SchemaTypesMetadata::default());
         let response = export_registry_types_ext(BrpJsonSchemaQueryFilter::default(), &world);
 
-        assert_eq!(response.definitions.len(), 5);
+        assert_eq!(
+            response.definitions.len(),
+            5,
+            "Expected 5 definitions, got: {:#?}",
+            response.definitions.keys()
+        );
         let response = export_registry_types_ext(
             BrpJsonSchemaQueryFilter {
                 without_crates: vec!["bevy_remote".to_string()],
diff --git a/crates/bevy_remote/src/lib.rs b/crates/bevy_remote/src/lib.rs
index 7d6a468796..c28240657c 100644
--- a/crates/bevy_remote/src/lib.rs
+++ b/crates/bevy_remote/src/lib.rs
@@ -517,6 +517,8 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::sync::RwLock;
 
+use crate::schemas::RegisterReflectJsonSchemas;
+
 pub mod builtin_methods;
 #[cfg(feature = "http")]
 pub mod http;
@@ -653,9 +655,7 @@ impl Default for RemotePlugin {
 
 impl Plugin for RemotePlugin {
     fn build(&self, app: &mut App) {
-        app.register_type::<schemas::open_rpc::OpenRpcDocument>()
-            .register_type_data::<schemas::open_rpc::OpenRpcDocument, schemas::ReflectJsonSchema>();
-
+        app.register_schema_base_types();
         let mut remote_methods = RemoteMethods::new();
 
         let plugin_methods = &mut *self.methods.write().unwrap();
diff --git a/crates/bevy_remote/src/schemas/mod.rs b/crates/bevy_remote/src/schemas/mod.rs
index fb8c6e788a..5970816d1b 100644
--- a/crates/bevy_remote/src/schemas/mod.rs
+++ b/crates/bevy_remote/src/schemas/mod.rs
@@ -12,7 +12,7 @@ use bevy_reflect::{
 };
 use core::any::TypeId;
 
-use crate::schemas::json_schema::JsonSchemaBevyType;
+use crate::schemas::{json_schema::JsonSchemaBevyType, open_rpc::OpenRpcDocument};
 
 pub mod json_schema;
 pub mod open_rpc;
@@ -37,6 +37,61 @@ impl<T: Reflect> FromType<T> for ReflectJsonSchemaForceAsArray {
     }
 }
 
+/// Helper trait
+pub(crate) trait RegisterReflectJsonSchemas {
+    /// Register types and or type data that are implemented by this crate
+    fn register_schema_base_types(&mut self) {
+        #[cfg(feature = "bevy_math")]
+        {
+            // self.register_type_data_internal::<bevy_math::Vec2, ReflectJsonSchemaForceAsArray>();
+            self.register_type_data_internal::<bevy_math::Vec3, ReflectJsonSchemaForceAsArray>();
+        }
+        self.register_type_internal::<OpenRpcDocument>();
+        self.register_type_data_internal::<OpenRpcDocument, ReflectJsonSchema>();
+    }
+    fn register_type_internal<T>(&mut self)
+    where
+        T: bevy_reflect::GetTypeRegistration;
+
+    fn register_type_data_internal<T, D>(&mut self)
+    where
+        T: Reflect + bevy_reflect::TypePath,
+        D: TypeData + FromType<T>;
+}
+
+impl RegisterReflectJsonSchemas for bevy_reflect::TypeRegistry {
+    fn register_type_data_internal<T, D>(&mut self)
+    where
+        T: Reflect + bevy_reflect::TypePath,
+        D: TypeData + FromType<T>,
+    {
+        self.register_type_data::<T, D>();
+    }
+
+    fn register_type_internal<T>(&mut self)
+    where
+        T: bevy_reflect::GetTypeRegistration,
+    {
+        self.register::<T>();
+    }
+}
+impl RegisterReflectJsonSchemas for bevy_app::App {
+    fn register_type_data_internal<T, D>(&mut self)
+    where
+        T: Reflect + bevy_reflect::TypePath,
+        D: TypeData + FromType<T>,
+    {
+        self.register_type_data::<T, D>();
+    }
+
+    fn register_type_internal<T>(&mut self)
+    where
+        T: bevy_reflect::GetTypeRegistration,
+    {
+        self.register_type::<T>();
+    }
+}
+
 /// Reflect-compatible custom JSON Schema for this type
 #[derive(Clone, Deref)]
 pub struct ReflectJsonSchema(pub JsonSchemaBevyType);