fix omit_exception to handle generators properly

Author: amirreza8002

django_valkey/base.py

@@ -48,14 +48,18 @@ def decorate(cls):
                     )
                 elif attr == "get_many":
                     setattr(cls, attr, decorator(attribute, return_value={}))
+                elif attr in {"iter_keys", "sscan_iter"}:
+                    setattr(cls, attr, decorator(attribute, gen=True))
                 else:
                     setattr(cls, attr, decorator(attribute))
         return cls
 
     return decorate
 
 
-def omit_exception(method: Callable | None = None, return_value: Any | None = None):
+def omit_exception(
+    method: Callable | None = None, return_value: Any | None = None, gen: bool = False
+):
     """
     Simple decorator that intercepts connection
     errors and ignores these if settings specify this.
@@ -102,10 +106,10 @@ async def _async_generator_decorator(self, *args, **kwargs):
         except ConnectionInterrupted as e:
             yield __handle_error(self, e)
 
-    # inspect.isfunction returns true for generators, so can't use that to check this
-    if not inspect.isasyncgenfunction(method) and not inspect.isgeneratorfunction(
-        method
-    ):
+    # sync generators (iter_keys, sscan_iter) are only generator on client class
+    # in the backend they are just a function (that returns a generator)
+    # so inspect.isgeneratorfunction does not work
+    if not inspect.isasyncgenfunction(method) and not gen:
         wrapper = _async_decorator if iscoroutinefunction(method) else _decorator
 
     # if method is a generator or async generator, it should be iterated over by this decorator