From 3696e53200a30f67b82a99da5978b568554908b7 Mon Sep 17 00:00:00 2001 From: Stijn de Gooijer Date: Sun, 19 Jan 2025 20:50:10 +0100 Subject: [PATCH] feat(python): Drop `nest-asyncio` in favor of custom logic (#20793) --- py-polars/polars/io/database/_utils.py | 46 ++++++++++++++----- py-polars/polars/meta/versions.py | 2 - py-polars/pyproject.toml | 3 +- py-polars/requirements-dev.txt | 1 - .../tests/unit/io/database/test_async.py | 5 +- 5 files changed, 36 insertions(+), 21 deletions(-) diff --git a/py-polars/polars/io/database/_utils.py b/py-polars/polars/io/database/_utils.py index 6ea01772cc98..30032c001e23 100644 --- a/py-polars/polars/io/database/_utils.py +++ b/py-polars/polars/io/database/_utils.py @@ -1,6 +1,9 @@ from __future__ import annotations +import asyncio import re +import threading +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any from polars.convert import from_arrow @@ -13,20 +16,39 @@ from polars._typing import SchemaDict -def _run_async(co: Coroutine[Any, Any, Any]) -> Any: - """Run asynchronous code as if it was synchronous.""" - import asyncio +def _run_async( + coroutine: Coroutine[Any, Any, Any], *, timeout: float | None = None +) -> Any: + """Run asynchronous code as if it were synchronous. - from polars._utils.unstable import issue_unstable_warning - from polars.dependencies import import_optional + This is required for execution in Jupyter notebook environments. + """ + # Implementation taken from StackOverflow answer here: + # https://stackoverflow.com/a/78911765/2344703 - issue_unstable_warning( - "Use of asynchronous connections is currently considered unstable " - "and unexpected issues may arise; if this happens, please report them." - ) - nest_asyncio = import_optional("nest_asyncio") - nest_asyncio.apply() - return asyncio.run(co) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # If there is no running loop, use `asyncio.run` normally + return asyncio.run(coroutine) + + def run_in_new_loop() -> Any: + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(coroutine) + finally: + new_loop.close() + + if threading.current_thread() is threading.main_thread(): + if not loop.is_running(): + return loop.run_until_complete(coroutine) + else: + with ThreadPoolExecutor() as pool: + future = pool.submit(run_in_new_loop) + return future.result(timeout=timeout) + else: + return asyncio.run_coroutine_threadsafe(coroutine, loop).result() def _read_sql_connectorx( diff --git a/py-polars/polars/meta/versions.py b/py-polars/polars/meta/versions.py index e7a7e3860e5c..efe73d8df9f1 100644 --- a/py-polars/polars/meta/versions.py +++ b/py-polars/polars/meta/versions.py @@ -30,7 +30,6 @@ def show_versions() -> None: fsspec: 2023.12.2 gevent: 24.2.1 matplotlib: 3.8.4 - nest_asyncio: 1.6.0 numpy: 1.26.4 openpyxl: 3.1.2 pandas: 2.2.2 @@ -85,7 +84,6 @@ def _get_dependency_list() -> list[str]: "google.auth", "great_tables", "matplotlib", - "nest_asyncio", "numpy", "openpyxl", "pandas", diff --git a/py-polars/pyproject.toml b/py-polars/pyproject.toml index 04cff909a76b..31145b724dc8 100644 --- a/py-polars/pyproject.toml +++ b/py-polars/pyproject.toml @@ -57,7 +57,7 @@ excel = ["polars[calamine,openpyxl,xlsx2csv,xlsxwriter]"] adbc = ["adbc-driver-manager[dbapi]", "adbc-driver-sqlite[dbapi]"] connectorx = ["connectorx >= 0.3.2"] sqlalchemy = ["sqlalchemy", "polars[pandas]"] -database = ["polars[adbc,connectorx,sqlalchemy]", "nest-asyncio"] +database = ["polars[adbc,connectorx,sqlalchemy]"] # Cloud fsspec = ["fsspec"] @@ -114,7 +114,6 @@ module = [ "kuzu", "matplotlib.*", "moto.server", - "nest_asyncio", "openpyxl", "polars.polars", "pyarrow.*", diff --git a/py-polars/requirements-dev.txt b/py-polars/requirements-dev.txt index 7fb358bced22..5ed6b1103d75 100644 --- a/py-polars/requirements-dev.txt +++ b/py-polars/requirements-dev.txt @@ -30,7 +30,6 @@ adbc-driver-sqlite; platform_system != 'Windows' aiosqlite connectorx kuzu -nest-asyncio # Cloud cloudpickle fsspec diff --git a/py-polars/tests/unit/io/database/test_async.py b/py-polars/tests/unit/io/database/test_async.py index 3bdea8207d2c..90354fc595d0 100644 --- a/py-polars/tests/unit/io/database/test_async.py +++ b/py-polars/tests/unit/io/database/test_async.py @@ -124,10 +124,7 @@ async def _nested_async_test(tmp_sqlite_db: Path) -> pl.DataFrame: reason="SQLAlchemy 2.0+ required for async tests", ) def test_read_async_nested(tmp_sqlite_db: Path) -> None: - # this tests validates that we can handle nested async calls. without - # the nested asyncio handling provided by `nest_asyncio` this test - # would raise a RuntimeError - + # This tests validates that we can handle nested async calls expected_frame = pl.DataFrame({"id": [1, 2], "name": ["misc", "other"]}) df = asyncio.run(_nested_async_test(tmp_sqlite_db)) assert_frame_equal(expected_frame, df)