ENH: Add Selectable (#305)

* ENH: Add Selectable

closes #304

* TST: Use mocking to get around sqlalchemy typing

Author: bashtage

pandas-stubs/io/sql.pyi

@@ -5,12 +5,14 @@ from typing import (
     Generator,
     Iterable,
     Literal,
+    Union,
     overload,
 )
 
 from pandas.core.base import PandasObject
 from pandas.core.frame import DataFrame
 import sqlalchemy.engine
+import sqlalchemy.sql.expression
 
 from pandas._typing import (
     DtypeArg,
@@ -20,10 +22,16 @@ from pandas._typing import (
 # TODO: Remove after switch to 1.5.x, moved to pandas.errors
 class DatabaseError(IOError): ...
 
+_SQLConnection = Union[
+    str,
+    sqlalchemy.engine.Connectable,
+    sqlite3.Connection,
+]
+
 @overload
 def read_sql_table(
     table_name: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    con: _SQLConnection,
     schema: str | None = ...,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
@@ -35,7 +43,7 @@ def read_sql_table(
 @overload
 def read_sql_table(
     table_name: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    con: _SQLConnection,
     schema: str | None = ...,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
@@ -45,8 +53,8 @@ def read_sql_table(
 ) -> DataFrame: ...
 @overload
 def read_sql_query(
-    sql: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    sql: str | sqlalchemy.sql.expression.Selectable,
+    con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -57,8 +65,8 @@ def read_sql_query(
 ) -> Generator[DataFrame, None, None]: ...
 @overload
 def read_sql_query(
-    sql: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    sql: str | sqlalchemy.sql.expression.Selectable,
+    con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -68,8 +76,8 @@ def read_sql_query(
 ) -> DataFrame: ...
 @overload
 def read_sql(
-    sql: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    sql: str | sqlalchemy.sql.expression.Selectable,
+    con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
@@ -80,8 +88,8 @@ def read_sql(
 ) -> Generator[DataFrame, None, None]: ...
 @overload
 def read_sql(
-    sql: str,
-    con: str | sqlalchemy.engine.Connectable | sqlite3.Connection,
+    sql: str | sqlalchemy.sql.expression.Selectable,
+    con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
     params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,

tests/test_io.py

@@ -42,6 +42,9 @@
 from pandas._testing import ensure_clean
 import pytest
 import sqlalchemy
+import sqlalchemy.ext.declarative
+import sqlalchemy.orm
+import sqlalchemy.orm.decl_api
 from typing_extensions import assert_type
 
 from tests import (
@@ -795,3 +798,24 @@ def test_csv_quoting():
             assert_type(DF.to_csv(path, quoting=csv.QUOTE_NONNUMERIC), None), type(None)
         )
         check(assert_type(DF.to_csv(path, quoting=csv.QUOTE_MINIMAL), None), type(None))
+
+
+def test_sqlalchemy_selectable() -> None:
+    with ensure_clean() as path:
+        db_uri = "sqlite:///" + path
+        engine = sqlalchemy.create_engine(db_uri)
+
+        if TYPE_CHECKING:
+            # Just type checking since underlying dB does not exist
+            class Base(metaclass=sqlalchemy.orm.decl_api.DeclarativeMeta):
+                __abstract__ = True
+
+            class Temp(Base):
+                __tablename__ = "part"
+                quantity = sqlalchemy.Column(sqlalchemy.Integer)
+
+            Session = sqlalchemy.orm.sessionmaker(engine)
+            with Session() as session:
+                pd.read_sql(
+                    session.query(Temp.quantity).statement, session.connection()
+                )