Accept sqlalchemy.engine.Engine for SQL IO API (read_sql, to_sql) (#281)

* Add sqlalchemy

* Add testcases using sqlalchemy engine and connection

* Accept sqlalchemy.engine.Engine

Author: wakabame

pandas-stubs/core/generic.pyi

@@ -160,7 +160,7 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
     def to_sql(
         self,
         name: _str,
-        con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+        con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
         schema: _str | None = ...,
         if_exists: Literal["fail", "replace", "append"] = ...,
         index: _bool = ...,

pandas-stubs/io/sql.pyi

@@ -23,7 +23,7 @@ class DatabaseError(IOError): ...
 @overload
 def read_sql_table(
     table_name: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     schema: str | None = ...,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
@@ -35,7 +35,7 @@ def read_sql_table(
 @overload
 def read_sql_table(
     table_name: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     schema: str | None = ...,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
@@ -46,7 +46,7 @@ def read_sql_table(
 @overload
 def read_sql_query(
     sql: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -58,7 +58,7 @@ def read_sql_query(
 @overload
 def read_sql_query(
     sql: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -69,7 +69,7 @@ def read_sql_query(
 @overload
 def read_sql(
     sql: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -81,7 +81,7 @@ def read_sql(
 @overload
 def read_sql(
     sql: str,
-    con: str | sqlalchemy.engine.Connection | sqlite3.Connection,
+    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,

pyproject.toml

@@ -57,6 +57,7 @@ odfpy = ">=1.4.1"
 xarray = ">=22.6.0"
 tabulate = ">=0.8.10"
 scipy = ">=1.9.1"
+SQLAlchemy = "^1.4.41"
 
 
 [build-system]

tests/test_io.py

@@ -43,6 +43,7 @@
 )
 from pandas._testing import ensure_clean
 import pytest
+import sqlalchemy
 from typing_extensions import assert_type
 
 from tests import check
@@ -692,6 +693,31 @@ def test_read_sql():
         con.close()
 
 
+def test_read_sql_via_sqlalchemy_connection():
+    with ensure_clean() as path:
+        db_uri = "sqlite:///" + path
+        engine = sqlalchemy.create_engine(db_uri)
+
+        with engine.connect() as conn:
+            check(assert_type(DF.to_sql("test", con=conn), Union[int, None]), int)
+            check(
+                assert_type(read_sql("select * from test", con=conn), DataFrame),
+                DataFrame,
+            )
+
+
+def test_read_sql_via_sqlalchemy_engine():
+    with ensure_clean() as path:
+        db_uri = "sqlite:///" + path
+        engine = sqlalchemy.create_engine(db_uri)
+
+        check(assert_type(DF.to_sql("test", con=engine), Union[int, None]), int)
+        check(
+            assert_type(read_sql("select * from test", con=engine), DataFrame),
+            DataFrame,
+        )
+
+
 def test_read_sql_generator():
     with ensure_clean() as path:
         con = sqlite3.connect(path)