Smarter

a_sync/a_sync/_helpers.py

@@ -6,7 +6,6 @@
 import a_sync.asyncio
 from a_sync import exceptions
 from a_sync._typing import *
-from a_sync.a_sync.function import ASyncFunction
 
 
 def _await(awaitable: Awaitable[T]) -> T:
@@ -15,9 +14,10 @@ def _await(awaitable: Awaitable[T]) -> T:
     except RuntimeError as e:
         if str(e) == "This event loop is already running":
             raise exceptions.SyncModeInAsyncContextError from None
-        raise e
+        raise
 
 def _asyncify(func: SyncFn[P, T], executor: Executor) -> CoroFn[P, T]:  # type: ignore [misc]
+    from a_sync.a_sync.function import ASyncFunction
     if asyncio.iscoroutinefunction(func) or isinstance(func, ASyncFunction):
         raise exceptions.FunctionNotSync(func)
     @functools.wraps(func)

a_sync/a_sync/function.py

@@ -89,25 +89,47 @@ def fn(self): # -> Union[SyncFn[[CoroFn[P, T]], MaybeAwaitable[T]], SyncFn[[Sync
         """Returns the final wrapped version of 'self._fn' decorated with all of the a_sync goodness."""
         return self._async_wrap if self._async_def else self._sync_wrap
 
-    def map(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> "TaskMapping[P, T]":
-        from a_sync import TaskMapping
-        return TaskMapping(self, *iterables, concurrency=concurrency, name=task_name, **kwargs)
-
-    async def any(self, *iterables: AnyIterable[object], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> bool:
-        return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **kwargs).any(pop=True, sync=False)
-
-    async def all(self, *iterables: AnyIterable[object], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> bool:
-        return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **kwargs).all(pop=True, sync=False)
-
-    async def min(self, *iterables: AnyIterable[object], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> T:
-        return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **kwargs).min(pop=True, sync=False)
-
-    async def max(self, *iterables: AnyIterable[object], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> T:
-        return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **kwargs).max(pop=True, sync=False)
-
-    async def sum(self, *iterables: AnyIterable[object], concurrency: Optional[int] = None, task_name: str = "", **kwargs: P.kwargs) -> T:
-        return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **kwargs).sum(pop=True, sync=False)
+    if sys.version_info >= (3, 11) or TYPE_CHECKING:
+        # we can specify P.args in python>=3.11 but in lower versions it causes a crash. Everything should still type check correctly on all versions.
+        def map(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> "TaskMapping[P, T]":
+            from a_sync import TaskMapping
+            return TaskMapping(self, *iterables, concurrency=concurrency, name=task_name, **function_kwargs)
+
+        async def any(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> bool:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).any(pop=True, sync=False)
+
+        async def all(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> bool:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).all(pop=True, sync=False)
+
+        async def min(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).min(pop=True, sync=False)
+
+        async def max(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).max(pop=True, sync=False)
+
+        async def sum(self, *iterables: AnyIterable[P.args], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).sum(pop=True, sync=False)
 
+    else:
+        def map(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> "TaskMapping[P, T]":
+            from a_sync import TaskMapping
+            return TaskMapping(self, *iterables, concurrency=concurrency, name=task_name, **function_kwargs)
+
+        async def any(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> bool:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).any(pop=True, sync=False)
+
+        async def all(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> bool:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).all(pop=True, sync=False)
+
+        async def min(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).min(pop=True, sync=False)
+
+        async def max(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).max(pop=True, sync=False)
+
+        async def sum(self, *iterables: AnyIterable[Any], concurrency: Optional[int] = None, task_name: str = "", **function_kwargs: P.kwargs) -> T:
+            return await self.map(*iterables, concurrency=concurrency, task_name=task_name, **function_kwargs).sum(pop=True, sync=False)
+
     @functools.cached_property
     def _sync_default(self) -> bool:
         """If user did not specify a default, we defer to the function. 'def' vs 'async def'"""

a_sync/asyncio/as_completed.py

@@ -137,6 +137,6 @@ async def __mapping_wrap(k: K, v: Awaitable[V], return_exceptions: bool = False)
     except Exception as e:
         if return_exceptions:
             return k, e
-        raise e
+        raise
 
 __all__ = ["as_completed", "as_completed_mapping"]

a_sync/asyncio/utils.py

@@ -6,7 +6,7 @@ def get_event_loop() -> asyncio.AbstractEventLoop:
         loop = asyncio.get_event_loop()
     except RuntimeError as e: # Necessary for use with multi-threaded applications.
         if not str(e).startswith("There is no current event loop in thread"):
-            raise e
+            raise
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
     return loop

a_sync/iter.py

@@ -170,12 +170,12 @@ def __cancel_cache_handle(self, instance: object) -> None:
         self._cache_handle.cancel()
 
 class _ASyncView(ASyncIterator[T]):
-    __aiterator__ = None
-    __iterator__ = None
+    __aiterator__: Optional[AsyncIterator[T]] = None
+    __iterator__: Optional[Iterator[T]] = None
     def __init__(
         self, 
         function: ViewFn[T], 
-        iterable: AsyncIterable[T],
+        iterable: AnyIterable[T],
     ) -> None:
         self._function = function
         self.__wrapped__ = iterable
@@ -184,7 +184,7 @@ def __init__(
         elif isinstance(iterable, Iterable):
             self.__iterator__ = iterable.__iter__()
         else:
-            raise TypeError(f"`iterable` must be AsyncIterable or Iterabe, you passed {iterable}")
+            raise TypeError(f"`iterable` must be AsyncIterable or Iterable, you passed {iterable}")
 
 @final  
 class ASyncFilter(_ASyncView[T]):
@@ -202,6 +202,8 @@ async def __anext__(self) -> T:
                         return obj
             except StopIteration:
                 pass
+        else:
+            raise TypeError(self.__wrapped__)
         raise StopAsyncIteration from None
     async def _check(self, obj: T) -> bool:
         checked = self._function(obj)

a_sync/primitives/queue.py

@@ -113,6 +113,7 @@ def __repr__(self) -> str:
         if self._unfinished_tasks:
             repr_string += f" pending={self._unfinished_tasks}"
         return f"{repr_string}>"
+    # NOTE: asyncio defines both this and __repr__
     def __str__(self) -> str:
         repr_string = f"<{type(self).__name__}"
         if self._name:
@@ -196,27 +197,25 @@ async def __worker_coro(self) -> NoReturn:
             while True:
                 try:
                     args, kwargs, fut = await self.get()
-                    if fut is None:
-                        # the weakref was already cleaned up, we don't need to process this item
-                        self.task_done()
-                        continue
-                    result = await self.func(*args, **kwargs)
-                    fut.set_result(result)
-                except asyncio.exceptions.InvalidStateError:
-                    logger.error("cannot set result for %s %s: %s", self.func.__name__, fut, result)
-                except Exception as e:
                     try:
-                        fut.set_exception(e)
+                        if fut is None:
+                            # the weakref was already cleaned up, we don't need to process this item
+                            self.task_done()
+                            continue
+                        result = await self.func(*args, **kwargs)
+                        fut.set_result(result)
                     except asyncio.exceptions.InvalidStateError:
-                        logger.error("cannot set exception for %s %s: %s", self.func.__name__, fut, e)
-                    except UnboundLocalError as u:
-                        logger.error("%s for %s is broken!!!", type(self).__name__, self.func)
-                        if str(e) != "local variable 'fut' referenced before assignment":
-                            logger.exception(u)
-                            raise u
-                        logger.exception(e)
-                        raise e
-                self.task_done()
+                        logger.error("cannot set result for %s %s: %s", self.func.__name__, fut, result)
+                    except Exception as e:
+                        try:
+                            fut.set_exception(e)
+                        except asyncio.exceptions.InvalidStateError:
+                            logger.error("cannot set exception for %s %s: %s", self.func.__name__, fut, e)
+                    self.task_done()
+                except Exception as e:
+                    logger.error("%s for %s is broken!!!", type(self).__name__, self.func)
+                    logger.exception(e)
+                    raise
 
 
 def _validate_args(i: int, can_return_less: bool) -> None:
@@ -337,13 +336,8 @@ async def __worker_coro(self) -> NoReturn:
                         fut.set_exception(e)
                     except asyncio.exceptions.InvalidStateError:
                         logger.error("cannot set exception for %s %s: %s", self.func.__name__, fut, e)
-                    except UnboundLocalError as u:
-                        if str(e) != "local variable 'fut' referenced before assignment":
-                            logger.exception(u)
-                            raise u
-                        raise e
                 self.task_done()
             except Exception as e:
                 logger.error("%s for %s is broken!!!", type(self).__name__, self.func)
                 logger.exception(e)
-                raise e
+                raise

a_sync/task.py

@@ -105,15 +105,12 @@ def __init__(
             @functools.wraps(wrapped_func)
             async def _wrapped_set_next(*args: P.args, __a_sync_recursion: int = 0, **kwargs: P.kwargs) -> V:
                 try:
-                    retval = await wrapped_func(*args, **kwargs)
-                    self._next.set()
-                    self._next.clear()
-                    return retval
+                    return await wrapped_func(*args, **kwargs)
                 except exceptions.SyncModeInAsyncContextError as e:
                     raise Exception(e, self.__wrapped__)
                 except TypeError as e:
                     if __a_sync_recursion > 2 or not (str(e).startswith(wrapped_func.__name__) and "got multiple values for argument" in str(e)):
-                        raise e
+                        raise
                     # NOTE: args ordering is clashing with provided kwargs. We can handle this in a hacky way.
                     # TODO: perform this check earlier and pre-prepare the args/kwargs ordering
                     new_args = list(args)
@@ -127,6 +124,9 @@ async def _wrapped_set_next(*args: P.args, __a_sync_recursion: int = 0, **kwargs
                         return await _wrapped_set_next(*new_args, **new_kwargs, __a_sync_recursion=__a_sync_recursion+1)
                     except TypeError as e2:
                         raise e.with_traceback(e.__traceback__) if str(e2) == "unsupported callable" else e2.with_traceback(e2.__traceback__)
+                finally:
+                    self._next.set()
+                    self._next.clear()
             self._wrapped_func = _wrapped_set_next
             init_loader_queue: Queue[Tuple[K, "asyncio.Future[V]"]] = Queue()
             self.__init_loader_coro = exhaust_iterator(self._tasks_for_iterables(*iterables), queue=init_loader_queue)

a_sync/utils/__init__.py

@@ -21,7 +21,7 @@ async def any(*awaitables) -> bool:
             if str(e) == "cannot reuse already awaited coroutine":
                 raise RuntimeError(str(e), fut) from e
             else:
-                raise e
+                raise
         if bool(result):
             for fut in futs:
                 fut.cancel()
@@ -37,7 +37,7 @@ async def all(*awaitables) -> bool:
             if str(e) == "cannot reuse already awaited coroutine":
                 raise RuntimeError(str(e), fut) from e
             else:
-                raise e
+                raise
         if not result:
             for fut in futs:
                 fut.cancel()

a_sync/utils/iterators.py

@@ -48,7 +48,8 @@ async def exhaust_iterators(iterators, *, queue: Optional[asyncio.Queue] = None,
     """
     for x in await asyncio.gather(*[exhaust_iterator(iterator, queue=queue) for iterator in iterators], return_exceptions=True):
         if isinstance(x, Exception):
-            raise x
+            # raise it with its original traceback instead of from here
+            raise x.with_traceback(x.__traceback__)
     if queue:
         queue.put_nowait(_Done())
         if join: