4 files changed, 43 insertions, 16 deletions

diff --git a/stockton-types/src/ent_map.rs b/stockton-types/src/ent_map.rs
index 213e481..e2eedcc 100644
--- a/stockton-types/src/ent_map.rs
+++ b/stockton-types/src/ent_map.rs
@@ -20,15 +20,30 @@ macro_rules! ent_map {
 		{
 			use stockton_bsp::lumps::entities::Entity as BSPEntity;	
 			use stockton_types::Entity;	
-			|ent: &BSPEntity| -> Box<dyn Entity> {
+			|ent: &BSPEntity| -> Option<Box<dyn Entity>> {
 				$(
 					if ent.attributes["classname"] == $name {
-						return Box::new($type {
-							$( $target : ent.attributes[$key].into() ),*
-						});
+						let mut valid = true;
+						{
+							$(let mut $target = false;);*
+							for key in ent.attributes.keys() {
+								$(if key == &$key {
+									$target = true;
+									continue;
+								});*
+							}
+							$(if !$target {
+								valid = false;
+							});*
+						}
+						if valid {
+							return Some(Box::new($type {
+								$( $target : ent.attributes[$key].into() ),*
+							}));
+						}
 					}
 				);*
-				panic!("Unrecognised Entity type: {:?}", ent);
+				None
 			}
 		}
 	}
diff --git a/stockton-types/src/world.rs b/stockton-types/src/world.rs
index 18a0edf..1196592 100644
--- a/stockton-types/src/world.rs
+++ b/stockton-types/src/world.rs
@@ -35,21 +35,20 @@ impl<'a> World<'a> {
 	///
 	/// `mapper` is called for each BSPEntity to map it to a concrete rust type.
 	pub fn new<F>(bsp: Pin<Box<BSPFile<'a>>>, mut mapper: F) -> Option<World<'a>>
-		where F: FnMut(&BSPEntity) -> (Box<Entity>, String) {
+		where F: FnMut(&BSPEntity) -> Option<(Box<Entity>, String)> {
 
 		let mut entities: Vec<(Box<Entity>, String)> = Vec::with_capacity(bsp.entities.entities.len());
 		for bsp_ent in bsp.entities.entities.iter() {
-			entities.push(mapper(&bsp_ent));
+			if let Some(result) = mapper(&bsp_ent) {
+				entities.push(result);
+			}
 		}
 
-		let store = EntityStore::from_entities(entities);
-		if store.is_none() {
-			return None;
-		}
+		let store = EntityStore::from_entities(entities)?;
 
 		Some(World {
 			map: bsp,
-			live_entities: store.unwrap()
+			live_entities: store
 		})
 	}
 }
\ No newline at end of file
diff --git a/stockton-types/tests/ent_map.rs b/stockton-types/tests/ent_map.rs
index 107620d..6c799ff 100644
--- a/stockton-types/tests/ent_map.rs
+++ b/stockton-types/tests/ent_map.rs
@@ -83,14 +83,14 @@ fn ent_map_macro() {
 		attributes: map![
 			"classname" => "A"
 		]
-	}).downcast_ref::<A>().unwrap(), &A);
+	}).unwrap().downcast_ref::<A>().unwrap(), &A);
 
 	assert_eq!(map_func(&BSPEntity {
 		attributes: map![
 			"classname" => "B",
 			"data" => "foobar"
 		]
-	}).downcast_ref::<B>().unwrap(), &B { data: "foobar".to_string() });
+	}).unwrap().downcast_ref::<B>().unwrap(), &B { data: "foobar".to_string() });
 
 	assert_eq!(map_func(&BSPEntity {
 		attributes: map![
@@ -98,5 +98,18 @@ fn ent_map_macro() {
 			"data" => "foobar",
 			"into" => "a b c"
 		]
-	}).downcast_ref::<C>().unwrap(), &C { data2: "foobar".to_string(), into: CustomStruct { one: 1, two: 2, three: 3 } });
+	}).unwrap().downcast_ref::<C>().unwrap(), &C { data2: "foobar".to_string(), into: CustomStruct { one: 1, two: 2, three: 3 } });
+
+	assert!(map_func(&BSPEntity {
+		attributes: map![
+			"classname" => "D"
+		]
+	}).is_none());
+
+	assert!(map_func(&BSPEntity {
+		attributes: map![
+			"classname" => "B",
+			"ebeb" => "foobar"
+		]
+	}).is_none());
 }
\ No newline at end of file
diff --git a/stockton-types/tests/world.rs b/stockton-types/tests/world.rs
index 83bb178..380918b 100644
--- a/stockton-types/tests/world.rs
+++ b/stockton-types/tests/world.rs
@@ -59,7 +59,7 @@ fn world_creation() {
 
 	let world = World::new(file, |ent: &BSPEntity| {
 		called_times += 1;
-		(Box::new(DummyEntity), ent.attributes.get("name").unwrap().clone().into())
+		Some((Box::new(DummyEntity), ent.attributes.get("name").unwrap().clone().into()))
 	}).unwrap();