diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index fac21c7c0ffc2..65887686d776d 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -5,7 +5,7 @@ use crate::{ self as bevy_ecs, bundle::Bundle, entity::{Entities, Entity}, - world::{FromWorld, World}, + world::{EntityMut, FromWorld, World}, }; use bevy_ecs_macros::SystemParam; use bevy_utils::tracing::{error, info}; @@ -835,6 +835,43 @@ impl<'w, 's, 'a> EntityCommands<'w, 's, 'a> { pub fn commands(&mut self) -> &mut Commands<'w, 's> { self.commands } + + /// Get mutable access to the target [`Entity`] through an [`EntityMut`]. + /// + /// # Panics + /// + /// The command will panic when applied if the associated entity does not exist. + /// + /// # Examples + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Component)] + /// # struct Health(u32); + /// # #[derive(Resource)] + /// # struct PlayerEntity { entity: Entity } + /// # fn my_system(mut commands: Commands, player: Res) { + /// commands + /// .entity(player.entity) + /// .for_entity_mut(|mut entity| { + /// entity.get_mut::().unwrap().0 *= 2; + /// }); + /// # } + /// # bevy_ecs::system::assert_is_system(my_system); + /// ``` + /// + /// # See also + /// + /// - [`add`][`Self::add`] for direct [`World`] access. + pub fn for_entity_mut( + &mut self, + f: impl for<'e> FnOnce(EntityMut<'e>) + Send + 'static, + ) -> &mut Self { + self.add(move |entity: Entity, world: &mut World| { + let entity = world.entity_mut(entity); + f(entity); + }) + } } impl Command for F @@ -1219,4 +1256,38 @@ mod tests { assert!(!world.contains_resource::>()); assert!(world.contains_resource::>()); } + + #[test] + fn modify_components() { + let mut world = World::default(); + + let mut command_queue = CommandQueue::default(); + + let entity = Commands::new(&mut command_queue, &world) + .spawn((W(1u32), W(2u64))) + .id(); + command_queue.apply(&mut world); + let results_before = world + .query::<(&W, &W)>() + .iter(&world) + .map(|(a, b)| (a.0, b.0)) + .collect::>(); + assert_eq!(results_before, vec![(1u32, 2u64)]); + + // test component modification + Commands::new(&mut command_queue, &world) + .entity(entity) + .for_entity_mut(|mut entity| { + entity.get_mut::>().unwrap().0 *= 2; + }); + + command_queue.apply(&mut world); + + let results_after = world + .query::<(&W, &W)>() + .iter(&world) + .map(|(a, b)| (a.0, b.0)) + .collect::>(); + assert_eq!(results_after, vec![(2u32, 2u64)]); + } }