Fix/context get local default (#1752)

* fix: correct default value return in context.get_local

* refactor: use correct get_local syntax

* lint: correct docstrings

---------

Co-authored-by: Kumaran Rajendhiran <[email protected]>

faststream/broker/middlewares/logging.py

@@ -36,7 +36,7 @@ async def on_consume(
         msg: "StreamMessage[Any]",
     ) -> "StreamMessage[Any]":
         if self.logger is not None:
-            c = context.get_local("log_context") or {}
+            c = context.get_local("log_context", {})
             self.logger.log(self.log_level, "Received", extra=c)
 
         return await super().on_consume(msg)
@@ -49,7 +49,7 @@ async def after_processed(
     ) -> bool:
         """Asynchronously called after processing."""
         if self.logger is not None:
-            c = context.get_local("log_context") or {}
+            c = context.get_local("log_context", {})
 
             if exc_type:
                 if issubclass(exc_type, IgnoredException):

faststream/log/logging.py

@@ -32,8 +32,8 @@ def __init__(
 
     def filter(self, record: LogRecord) -> bool:
         if is_suitable := super().filter(record):
-            log_context: Mapping[str, str] = (
-                context.get_local("log_context") or self.default_context
+            log_context: Mapping[str, str] = context.get_local(
+                "log_context", self.default_context
             )
 
             for k, v in log_context.items():

faststream/nats/broker/registrator.py

@@ -26,7 +26,7 @@
 
 
 class NatsRegistrator(ABCBroker["Msg"]):
-    """Includable to RabbitBroker router."""
+    """Includable to NatsBroker router."""
 
     _subscribers: Dict[int, "AsyncAPISubscriber"]
     _publishers: Dict[int, "AsyncAPIPublisher"]

faststream/redis/broker/registrator.py

@@ -25,7 +25,7 @@
 
 
 class RedisRegistrator(ABCBroker[UnifyRedisDict]):
-    """Includable to RabbitBroker router."""
+    """Includable to RedisBroker router."""
 
     _subscribers: Dict[int, "SubsciberType"]
     _publishers: Dict[int, "PublisherType"]

faststream/utils/context/repository.py

@@ -66,7 +66,7 @@ def set_local(self, key: str, value: Any) -> "Token[Any]":
         """
         context_var = self._scope_context.get(key)
         if context_var is None:
-            context_var = ContextVar(key, default=None)
+            context_var = ContextVar(key, default=EMPTY)
             self._scope_context[key] = context_var
         return context_var.set(value)
 
@@ -92,10 +92,12 @@ def get_local(self, key: str, default: Any = None) -> Any:
         Returns:
             The value of the local variable.
         """
-        if (context_var := self._scope_context.get(key)) is not None:
-            return context_var.get()
-        else:
-            return default
+        value = default
+        if (context_var := self._scope_context.get(key)) is not None and (
+            context_value := context_var.get()
+        ) is not EMPTY:
+            value = context_value
+        return value
 
     @contextmanager
     def scope(self, key: str, value: Any) -> Iterator[None]:

pyproject.toml

@@ -291,6 +291,7 @@ minversion = "7.0"
 addopts = "-q -m 'not slow'"
 testpaths = ["tests"]
 markers = ["rabbit", "kafka", "confluent", "nats", "redis", "slow", "all"]
+asyncio_default_fixture_loop_scope = "function"
 
 [tool.coverage.run]
 parallel = true

tests/utils/context/test_main.py

@@ -154,6 +154,15 @@ def use(
         use()
 
 
+def test_local_default(context: ContextRepo):
+    key = "some-key"
+
+    tag = context.set_local(key, "useless")
+    context.reset_local(key, tag)
+
+    assert context.get_local(key, 1) == 1
+
+
 def test_initial():
     @apply_types
     def use(