From 02f9f375c119b487ede481cce0a2bd27f6fb3fda Mon Sep 17 00:00:00 2001
From: Freja Roberts <freja.mroberts@gmail.com>
Date: Sun, 11 Aug 2024 18:46:50 +0200
Subject: [PATCH] feat: expect

---
 src/fetch/expect.rs | 46 +++++++++++++++++++++++++++++++++++++++++++++
 src/fetch/ext.rs    |  8 ++++++++
 src/fetch/mod.rs    |  1 +
 tests/expect.rs     | 12 ++++++++++++
 4 files changed, 67 insertions(+)
 create mode 100644 src/fetch/expect.rs
 create mode 100644 tests/expect.rs

diff --git a/src/fetch/expect.rs b/src/fetch/expect.rs
new file mode 100644
index 0000000..f00188f
--- /dev/null
+++ b/src/fetch/expect.rs
@@ -0,0 +1,46 @@
+use std::borrow::Cow;
+
+use super::{Fetch, FetchItem};
+
+/// Expect the query to match, panic otherwise
+pub struct Expect<Q> {
+    msg: Cow<'static, str>,
+    fetch: Q,
+}
+
+impl<Q> Expect<Q> {
+    /// Expect the query to match, panic otherwise
+    pub fn new(fetch: Q, msg: impl Into<Cow<'static, str>>) -> Self {
+        Self {
+            fetch,
+            msg: msg.into(),
+        }
+    }
+}
+
+impl<'q, Q: FetchItem<'q>> FetchItem<'q> for Expect<Q> {
+    type Item = Q::Item;
+}
+
+impl<'w, Q: Fetch<'w>> Fetch<'w> for Expect<Q> {
+    const MUTABLE: bool = Q::MUTABLE;
+
+    type Prepared = Q::Prepared;
+
+    fn prepare(&'w self, data: super::FetchPrepareData<'w>) -> Option<Self::Prepared> {
+        Some(self.fetch.prepare(data).expect(&self.msg))
+    }
+
+    fn filter_arch(&self, _: super::FetchAccessData) -> bool {
+        true
+    }
+
+    fn access(&self, data: super::FetchAccessData, dst: &mut Vec<crate::system::Access>) {
+        self.fetch.access(data, dst)
+    }
+
+    fn describe(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_str("expect ")?;
+        self.fetch.describe(f)
+    }
+}
diff --git a/src/fetch/ext.rs b/src/fetch/ext.rs
index c33a3b3..a00ba44 100644
--- a/src/fetch/ext.rs
+++ b/src/fetch/ext.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use crate::{
     component::ComponentValue,
     filter::{Cmp, Equal, Filtered, Greater, GreaterEq, Less, LessEq},
@@ -9,6 +11,7 @@ use super::{
     as_deref::AsDeref,
     cloned::Cloned,
     copied::Copied,
+    expect::Expect,
     opt::{Opt, OptOr},
     source::{FetchSource, FromRelation, Traverse},
     transform::Added,
@@ -208,6 +211,11 @@ pub trait FetchExt: Sized {
     {
         Filtered::new(self, filter, true)
     }
+
+    /// Expect the query to match, panic otherwise
+    fn expect(self, msg: impl Into<Cow<'static, str>>) -> Expect<Self> {
+        Expect::new(self, msg.into())
+    }
 }
 
 impl<F> FetchExt for F where F: for<'x> Fetch<'x> {}
diff --git a/src/fetch/mod.rs b/src/fetch/mod.rs
index 9179983..7a5b571 100644
--- a/src/fetch/mod.rs
+++ b/src/fetch/mod.rs
@@ -4,6 +4,7 @@ mod component;
 mod component_mut;
 mod copied;
 mod entity_ref;
+pub mod expect;
 mod ext;
 mod map;
 mod maybe_mut;
diff --git a/tests/expect.rs b/tests/expect.rs
new file mode 100644
index 0000000..98393d8
--- /dev/null
+++ b/tests/expect.rs
@@ -0,0 +1,12 @@
+use flax::{components::name, Entity, FetchExt, Query, World};
+
+#[test]
+#[should_panic(expected = "name must be present")]
+fn expect_present() {
+    let mut world = World::new();
+
+    Entity::builder().spawn(&mut world);
+    let mut query = Query::new(name().cloned().expect("name must be present"));
+
+    let _ = query.collect_vec(&world);
+}