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 { + msg: Cow<'static, str>, + fetch: Q, +} + +impl Expect { + /// Expect the query to match, panic otherwise + pub fn new(fetch: Q, msg: impl Into>) -> Self { + Self { + fetch, + msg: msg.into(), + } + } +} + +impl<'q, Q: FetchItem<'q>> FetchItem<'q> for Expect { + type Item = Q::Item; +} + +impl<'w, Q: Fetch<'w>> Fetch<'w> for Expect { + const MUTABLE: bool = Q::MUTABLE; + + type Prepared = Q::Prepared; + + fn prepare(&'w self, data: super::FetchPrepareData<'w>) -> Option { + 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) { + 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>) -> Expect { + Expect::new(self, msg.into()) + } } impl 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); +}