From caf3d8b3ef70ab861de7056afb29c70df1a9e129 Mon Sep 17 00:00:00 2001
From: Carter Anderson <mcanders1@gmail.com>
Date: Sat, 19 Dec 2020 10:31:47 -0800
Subject: [PATCH] add with_enter_stage (and other variants) (#1091)

---
 crates/bevy_ecs/src/schedule/state.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/crates/bevy_ecs/src/schedule/state.rs b/crates/bevy_ecs/src/schedule/state.rs
index 1fa6f0b2d9..f54a5f0eba 100644
--- a/crates/bevy_ecs/src/schedule/state.rs
+++ b/crates/bevy_ecs/src/schedule/state.rs
@@ -33,6 +33,21 @@ impl<T> Default for StateStage<T> {
 
 #[allow(clippy::mem_discriminant_non_enum)]
 impl<T> StateStage<T> {
+    pub fn with_enter_stage<S: Stage>(mut self, state: T, stage: S) -> Self {
+        self.set_enter_stage(state, stage);
+        self
+    }
+
+    pub fn with_exit_stage<S: Stage>(mut self, state: T, stage: S) -> Self {
+        self.set_exit_stage(state, stage);
+        self
+    }
+
+    pub fn with_update_stage<S: Stage>(mut self, state: T, stage: S) -> Self {
+        self.set_update_stage(state, stage);
+        self
+    }
+
     pub fn set_enter_stage<S: Stage>(&mut self, state: T, stage: S) -> &mut Self {
         let stages = self.state_stages(state);
         stages.enter = Box::new(stage);