From f9853cbbc2d52c1fe07b0ad0829f9cec22de0f9c Mon Sep 17 00:00:00 2001
From: James O'Brien <james.obrien@drafly.net>
Date: Thu, 1 Sep 2022 22:06:46 +0000
Subject: [PATCH] Add get_entity to Commands (#5854)

# Objective

- Fixes #5850

## Solution

- As described in the issue, added a `get_entity` method on `Commands` that returns an `Option<EntityCommands>`

## Changelog
- Added the new method with a simple doc test
- I have re-used `get_entity` in `entity`, similarly to how `get_single` is used in `single` while additionally preserving the error message
- Add `#[inline]` to both functions

Entities that have commands queued to despawn system will still return commands when `get_entity` is called but that is representative of the fact that the entity is still around until those commands are flushed.

A potential `contains_entity` could also be added in this PR if desired, that would effectively be replacing Entities.contains but may be more discoverable if this is a common use case.


Co-authored-by: Carter Anderson <mcanders1@gmail.com>
---
 crates/bevy_ecs/src/system/commands/mod.rs | 44 ++++++++++++++++++----
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs
index 19c057bc4b664..1054538b19c84 100644
--- a/crates/bevy_ecs/src/system/commands/mod.rs
+++ b/crates/bevy_ecs/src/system/commands/mod.rs
@@ -243,17 +243,47 @@ impl<'w, 's> Commands<'w, 's> {
     /// }
     /// # bevy_ecs::system::assert_is_system(example_system);
     /// ```
+    #[inline]
     #[track_caller]
     pub fn entity<'a>(&'a mut self, entity: Entity) -> EntityCommands<'w, 's, 'a> {
-        assert!(
-            self.entities.contains(entity),
-            "Attempting to create an EntityCommands for entity {:?}, which doesn't exist.",
-            entity
-        );
-        EntityCommands {
+        self.get_entity(entity).unwrap_or_else(|| {
+            panic!(
+                "Attempting to create an EntityCommands for entity {:?}, which doesn't exist.",
+                entity
+            )
+        })
+    }
+
+    /// Returns an option containing an [`EntityCommands`] builder for the requested [`Entity`] if it exists, otherwise `None`.
+    /// This does not ensure that the commands will succeed as the entity may no longer exist by the time the associated commands are executed.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    ///
+    /// #[derive(Component)]
+    /// struct Label(&'static str);
+
+    /// fn example_system(mut commands: Commands) {
+    ///     // Create a new, empty entity
+    ///     let entity = commands.spawn().id();
+    ///
+    ///     // Get the entity if it still exists, which it will in this case
+    ///     if let Some(mut entity_commands) = commands.get_entity(entity) {
+    ///         // adds a single component to the entity
+    ///         entity_commands.insert(Label("hello world"));
+    ///     }      
+    /// }
+    /// # bevy_ecs::system::assert_is_system(example_system);
+    /// ```
+    #[inline]
+    #[track_caller]
+    pub fn get_entity<'a>(&'a mut self, entity: Entity) -> Option<EntityCommands<'w, 's, 'a>> {
+        self.entities.contains(entity).then_some(EntityCommands {
             entity,
             commands: self,
-        }
+        })
     }
 
     /// Spawns entities to the [`World`] according to the given iterator (or a type that can