fixing mypy issues

qweeze · Apr 8, 2024 · 52dc572 · 52dc572
1 parent c52ad08
commit 52dc572
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,6 +27,8 @@ mmh3 = "^4.0.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.0"
+requests = "^2.31.0"
+types-requests = "^2.31.0.20240406"
 
 [tool.black]
 line-length = 110

diff --git a/rstream/consumer.py b/rstream/consumer.py
@@ -454,10 +454,12 @@ async def reconnect_stream(self, stream: str, offset: Optional[int] = None) -> N
             if stream == self._subscribers[subscriber_id].stream:
                 curr_subscriber = self._subscribers[subscriber_id]
                 curr_subscriber_id = subscriber_id
-        del self._subscribers[curr_subscriber_id]
+        if curr_subscriber_id is not None:
+            del self._subscribers[curr_subscriber_id]
 
         if stream in self._clients:
-            await self._clients[stream].free_available_id(curr_subscriber.subscription_id)
+            if curr_subscriber is not None:
+                await self._clients[stream].free_available_id(curr_subscriber.subscription_id)
             await self._clients[stream].close()
             del self._clients[stream]
 
@@ -519,7 +521,8 @@ async def _maybe_clean_up_during_lost_connection(self, stream: str):
 
         if stream in self._clients:
             await self._clients[stream].remove_stream(stream)
-            await self._clients[stream].free_available_id(curr_subscriber.subscription_id)
+            if curr_subscriber is not None:
+                await self._clients[stream].free_available_id(curr_subscriber.subscription_id)
             if await self._clients[stream].get_stream_count() == 0:
                 await self._clients[stream].close()
             del self._clients[stream]