Explicit error for query:next

Ukendio · Jul 30, 2024 · d7bda01 · d7bda01
1 parent d5baf52
commit d7bda01
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/src/init.luau b/src/init.luau
@@ -813,7 +813,7 @@ do
         end
     end
 
-    local function world_query_next(): any
+    local function world_query_iter_next(): any
         local entityId = entities[i]
   		while entityId == nil do
  			lastArchetype += 1
@@ -1048,9 +1048,17 @@ do
         if not drain then
             query_init(query)
         end
-        return world_query_next
+        return world_query_iter_next
     end
 
+    local function world_query_next()
+        if not drain then
+            error("Did you forget to call query:drain()?")
+        end
+        return world_query_iter_next()
+    end
+
+
     local it = {
         __iter = world_query_iter,
         drain = world_query_drain,

diff --git a/test/tests.luau b/test/tests.luau
@@ -296,14 +296,19 @@ TEST("world:query()", function()
   		world:set(eAB, A, true)
   		world:set(eAB, B, true)
 
-  		local q = world:query(A)
+  		local q = world:query(A):drain()
 
-  		local e, data = q:next()
+  		local e, data = q.next()
   		while e do
-            if e ~= eA and e ~= eAB then
+            if e == eA then
+                CHECK(data)
+            elseif e == eAB then
+                CHECK(data)
+            else
                 CHECK(false)
             end
- 			e, data = q:next()
+
+ 			e, data = q.next()
   		end
         CHECK(true)
    	end
@@ -743,7 +748,7 @@ do
         added = true
         local q = world:query(T):without(PreviousT):drain()
         return function()
-            local id, data = q:next()
+            local id, data = q.next()
             if not id then
                 return nil
             end
@@ -762,7 +767,7 @@ do
         local q = world:query(T, PreviousT):drain()
 
         return function()
-            local id, new, old = q:next()
+            local id, new, old = q.next()
             while true do
                 if not id then
                     return nil
@@ -788,7 +793,7 @@ do
 
         local q = world:query(PreviousT):without(T):drain()
         return function()
-            local id = q:next()
+            local id = q.next()
             if id then
                 world:remove(id, PreviousT)
             end