Fix sql params type

pandas-dev · Jul 25, 2023 · 3355333 · 3355333
1 parent 89e083b
commit 3355333
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 4 deletions.
diff --git a/pandas-stubs/io/sql.pyi b/pandas-stubs/io/sql.pyi
@@ -2,6 +2,7 @@ from collections.abc import (
     Callable,
     Generator,
     Iterable,
+    Mapping,
 )
 import sqlite3
 from typing import (
@@ -65,7 +66,7 @@ def read_sql_query(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
+    params: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     *,
     chunksize: int,
@@ -78,7 +79,7 @@ def read_sql_query(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
+    params: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     chunksize: None = ...,
     dtype: DtypeArg | None = ...,
@@ -90,7 +91,7 @@ def read_sql(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
+    params: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     columns: list[str] = ...,
     *,
@@ -104,7 +105,7 @@ def read_sql(
     con: _SQLConnection,
     index_col: str | list[str] | None = ...,
     coerce_float: bool = ...,
-    params: list[str] | tuple[str, ...] | dict[str, str] | None = ...,
+    params: list[Any] | tuple[Any, ...] | Mapping[str, Any] | None = ...,
     parse_dates: list[str] | dict[str, str] | dict[str, dict[str, Any]] | None = ...,
     columns: list[str] = ...,
     chunksize: None = ...,

diff --git a/tests/test_io.py b/tests/test_io.py
@@ -1141,6 +1141,26 @@ def test_read_sql_via_sqlalchemy_engine():
         engine.dispose()
 
 
+def test_read_sql_via_sqlalchemy_engine_with_params():
+    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 where a = :a and b = :b",
+                    con=engine,
+                    params={"a": 2, "b": 0.0},
+                ),
+                DataFrame,
+            ),
+            DataFrame,
+        )
+        engine.dispose()
+
+
 def test_read_sql_generator():
     with ensure_clean() as path:
         con = sqlite3.connect(path)
@@ -1200,6 +1220,26 @@ def test_read_sql_query_generator():
         con.close()
 
 
+def test_read_sql_query_via_sqlalchemy_engine_with_params():
+    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_query(
+                    "select * from test where a = :a and b = :b",
+                    con=engine,
+                    params={"a": 2, "b": 0.0},
+                ),
+                DataFrame,
+            ),
+            DataFrame,
+        )
+        engine.dispose()
+
+
 def test_read_html():
     check(assert_type(DF.to_html(), str), str)
     with ensure_clean() as path: