From 692027413feea39c13bc22020843c3a6c7e329b1 Mon Sep 17 00:00:00 2001
From: Aidan Holland <aidan@censys.com>
Date: Fri, 13 Dec 2024 16:56:27 -0500
Subject: [PATCH] chore: Added total parameter to the Query object

---
 censys/search/v2/api.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/censys/search/v2/api.py b/censys/search/v2/api.py
index 3da9359b..1edc2ddb 100644
--- a/censys/search/v2/api.py
+++ b/censys/search/v2/api.py
@@ -100,6 +100,9 @@ class Query(Iterable):
         For more details, see our documentation: https://search.censys.io/api
         """
 
+        # Total number of results (Set after first query)
+        total: Optional[int] = None
+
         def __init__(
             self,
             api: "CensysSearchAPIv2",
@@ -162,8 +165,9 @@ def __call__(self, per_page: Optional[int] = None) -> List[dict]:
             )
             self.page += 1
             result = payload["result"]
+            self.total = result["total"]
             self.nextCursor = result["links"].get("next")
-            if result["total"] == 0 or not self.nextCursor:
+            if self.total == 0 or not self.nextCursor:
                 self.pages = 0
             return result["hits"]