add with_enter_stage (and other variants) (bevyengine#1091)

DJMcNab · Dec 19, 2020 · caf3d8b · caf3d8b
1 parent 9db38b7
commit caf3d8b
Showing 1 changed file with 15 additions and 0 deletions.
diff --git 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);