feat: expect

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)
+    }
+}

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> {}

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;

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);
+}