mypy: check untyped defs

nats/aio/subscription.py

@@ -84,9 +84,9 @@ def __init__(
         self._pending_queue: asyncio.Queue[Msg] = asyncio.Queue(
             maxsize=pending_msgs_limit
         )
-        self._pending_size = 0
-        self._wait_for_msgs_task = None
-        self._message_iterator = None
+        self._pending_size: int = 0
+        self._wait_for_msgs_task: asyncio.Task | None = None
+        self._message_iterator: _SubscriptionMessageIterator | None = None
 
         # For JetStream enabled subscriptions.
         self._jsi: JetStreamContext._JSI | None = None
@@ -182,8 +182,10 @@ def _start(self, error_cb):
         Creates the resources for the subscription to start processing messages.
         """
         if self._cb:
-            if not asyncio.iscoroutinefunction(self._cb) and \
-                    not (hasattr(self._cb, "func") and asyncio.iscoroutinefunction(self._cb.func)):
+            # TODO(@orsinium): What's this check? Does `self._cb` ever have `func`?
+            wrapped = getattr(self._cb, "func", None)
+            if (not asyncio.iscoroutinefunction(self._cb)
+                    and not asyncio.iscoroutinefunction(wrapped)):
                 raise errors.Error(
                     "nats: must use coroutine for subscriptions"
                 )
@@ -302,7 +304,8 @@ async def _wait_for_msgs(self, error_cb) -> None:
                     self._pending_queue.task_done()
 
                 # Apply auto unsubscribe checks after having processed last msg.
-                if self._max_msgs > 0 and self._received >= self._max_msgs and self._pending_queue.empty:
+                if self._max_msgs > 0 and self._received >= self._max_msgs and self._pending_queue.empty(
+                ):
                     self._stop_processing()
 
             except asyncio.CancelledError:

nats/js/kv.py

@@ -24,6 +24,7 @@
 
 if TYPE_CHECKING:
     from nats.js import JetStreamContext
+    from nats.aio.msg import Msg
 
 KV_OP = "KV-Operation"
 KV_DEL = "DEL"
@@ -73,7 +74,7 @@ class Entry:
         value: bytes | None
         revision: int | None
         delta: int | None
-        created: int | None
+        created: datetime.datetime | None
         operation: str | None
 
     @dataclass(frozen=True)
@@ -294,12 +295,17 @@ async def status(self) -> BucketStatus:
         return KeyValue.BucketStatus(stream_info=info, bucket=self._name)
 
     class KeyWatcher:
+        _js: JetStreamContext
+        _updates: asyncio.Queue[KeyValue.Entry | None]
+        _sub: JetStreamContext.PushSubscription | None
+        _pending: int | None
+        _init_done: bool
 
-        def __init__(self, js):
+        def __init__(self, js: JetStreamContext):
             self._js = js
             self._updates = asyncio.Queue(maxsize=256)
             self._sub = None
-            self._pending: int | None = None
+            self._pending = None
 
             # init done means that the nil marker has been sent,
             # once this is sent it won't be sent anymore.
@@ -309,7 +315,8 @@ async def stop(self):
             """
             stop will stop this watcher.
             """
-            await self._sub.unsubscribe()
+            if self._sub is not None:
+                await self._sub.unsubscribe()
 
         async def updates(self, timeout=5):
             """
@@ -394,10 +401,10 @@ async def watch(
         there are no pending updates.
         """
         subject = f"{self._pre}{keys}"
-        watcher = KeyValue.KeyWatcher(self)
+        watcher = KeyValue.KeyWatcher(self._js)
         init_setup: asyncio.Future[bool] = asyncio.Future()
 
-        async def watch_updates(msg):
+        async def watch_updates(msg: Msg) -> None:
             if not init_setup.done():
                 await asyncio.wait_for(init_setup, timeout=self._js._timeout)
 
@@ -418,7 +425,8 @@ async def watch_updates(msg):
                 bucket=self._name,
                 key=msg.subject[len(self._pre):],
                 value=msg.data,
-                revision=meta.sequence.stream,
+                revision=meta.sequence.stream
+                if meta.sequence is not None else None,
                 delta=meta.num_pending,
                 created=meta.timestamp,
                 operation=op,