Create a Dynamic Scene From a Query Filter

crates/bevy_scene/src/dynamic_scene.rs

@@ -2,7 +2,8 @@ use crate::{serde::SceneSerializer, Scene, SceneSpawnError};
 use anyhow::Result;
 use bevy_app::AppTypeRegistry;
 use bevy_ecs::{
-    entity::EntityMap,
+    entity::{Entity, EntityMap},
+    query::ReadOnlyWorldQuery,
     reflect::{ReflectComponent, ReflectMapEntities},
     world::World,
 };
@@ -37,6 +38,84 @@ impl DynamicScene {
         Self::from_world(&scene.world, type_registry)
     }
 
+    /// Create a Dynamic Scene with specific entities.
+    ///
+    /// The generic parameter is used as a [`Query`](bevy_ecs::system::Query) filter.
+    ///
+    /// The created scene will include only the entities that match the query
+    /// filter provided. All components that impl `Reflect` will be included.
+    ///
+    /// # Example
+    ///
+    /// Here, the function is creating a scene with all the componentA's in the world, but none of the componentB's.
+    ///
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    /// use bevy_scene::DynamicScene;
+    /// use bevy_app::AppTypeRegistry;
+    ///
+    /// #[derive(Component)]
+    /// struct ComponentA;
+    ///
+    /// #[derive(Component)]
+    /// struct ComponentB;
+    ///
+    /// let mut world = World::new();
+    /// world.init_resource::<AppTypeRegistry>();
+    /// let type_registry = world.resource::<AppTypeRegistry>().clone();
+    ///
+    /// let my_scene = DynamicScene::from_query_filter::<(With<ComponentA>, Without<ComponentB>)>(
+    ///   &mut world,
+    ///   &type_registry,
+    /// );
+    /// ```
+    pub fn from_query_filter<F>(world: &mut World, type_registry: &AppTypeRegistry) -> Self
+    where
+        F: ReadOnlyWorldQuery + 'static,
+    {
+        let mut query = world.query_filtered::<Entity, F>();
+
+        let type_registry = type_registry.read();
+
+        let entities = query
+            .iter(world)
+            .map(|entity| {
+                let get_reflect_by_id = |id| {
+                    world
+                        .components()
+                        .get_info(id)
+                        .and_then(|info| {
+                            type_registry.get(
+                                info.type_id()
+                                    .expect("Cannot find the type id for the component."),
+                            )
+                        })
+                        .and_then(|reg| reg.data::<ReflectComponent>())
+                        .and_then(|rc| rc.reflect(world, entity))
+                        .map(|c| c.clone_value())
+                };
+
+                // Components contains all the reflectable components
+                // per entity that satisfies the query.
+                let components = world
+                    .entities()
+                    .get(entity)
+                    .and_then(|eloc| world.archetypes().get(eloc.archetype_id))
+                    .into_iter()
+                    .flat_map(|a| a.components())
+                    .filter_map(get_reflect_by_id)
+                    .collect();
+
+                DynamicEntity {
+                    entity: entity.id(),
+                    components,
+                }
+            })
+            .collect();
+
+        DynamicScene { entities }
+    }
+
     /// Create a new dynamic scene from a given world.
     pub fn from_world(world: &World, type_registry: &TypeRegistryArc) -> Self {
         let mut scene = DynamicScene::default();
@@ -145,3 +224,33 @@ where
         .new_line("\n".to_string());
     ron::ser::to_string_pretty(&serialize, pretty_config)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::DynamicScene;
+    use bevy_app::AppTypeRegistry;
+    use bevy_ecs::prelude::*;
+
+    #[test]
+    fn from_query_filter_test() {
+        #[derive(Component)]
+        struct ComponentA;
+
+        #[derive(Component)]
+        struct ComponentB;
+
+        let mut world = World::new();
+        world.init_resource::<AppTypeRegistry>();
+        let type_registry = world.resource::<AppTypeRegistry>().clone();
+
+        let _entity1 = world.spawn().insert(ComponentA);
+        let _entity2 = world.spawn().insert(ComponentB);
+
+        let my_scene = DynamicScene::from_query_filter::<(With<ComponentA>, Without<ComponentB>)>(
+            &mut world,
+            &type_registry,
+        );
+
+        assert_eq!(my_scene.entities.len(), 1);
+    }
+}