From fdac05cc10804502c89fe09644ddf5b455d2a647 Mon Sep 17 00:00:00 2001 From: Andrew Wieteska Date: Thu, 17 Dec 2020 02:24:35 -0500 Subject: [PATCH 1/3] typing --- pandas/io/sql.py | 126 +++++++++++++++++++++++++++++------------------ 1 file changed, 79 insertions(+), 47 deletions(-) diff --git a/pandas/io/sql.py b/pandas/io/sql.py index 23f992ceb009a..66c4fbee3f1c9 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -7,7 +7,7 @@ from datetime import date, datetime, time from functools import partial import re -from typing import Iterator, List, Optional, Union, overload +from typing import Iterator, List, Optional, Sequence, Union, overload import warnings import numpy as np @@ -77,7 +77,7 @@ def _process_parse_dates_argument(parse_dates): return parse_dates -def _handle_date_column(col, utc=None, format=None): +def _handle_date_column(col, utc: Optional[bool] = None, format: Optional[str] = None): if isinstance(format, dict): # GH35185 Allow custom error values in parse_dates argument of # read_sql like functions. @@ -124,7 +124,13 @@ def _parse_date_columns(data_frame, parse_dates): return data_frame -def _wrap_result(data, columns, index_col=None, coerce_float=True, parse_dates=None): +def _wrap_result( + data, + columns, + index_col=None, + coerce_float: bool = True, + parse_dates: Optional[bool] = None, +): """Wrap result set of query in a DataFrame.""" frame = DataFrame.from_records(data, columns=columns, coerce_float=coerce_float) @@ -197,11 +203,11 @@ def read_sql_table( def read_sql_table( - table_name, + table_name: str, con, - schema=None, - index_col=None, - coerce_float=True, + schema: Optional[str] = None, + index_col: Optional[Union[str, Sequence[str]]] = None, + coerce_float: bool = True, parse_dates=None, columns=None, chunksize: Optional[int] = None, @@ -321,7 +327,7 @@ def read_sql_query( sql, con, index_col=None, - coerce_float=True, + coerce_float: bool = True, params=None, parse_dates=None, chunksize: Optional[int] = None, @@ -420,8 +426,8 @@ def read_sql( def read_sql( sql, con, - index_col=None, - coerce_float=True, + index_col: Optional[Union[str, Sequence[str]]] = None, + coerce_float: bool = True, params=None, parse_dates=None, columns=None, @@ -582,15 +588,15 @@ def read_sql( def to_sql( frame, - name, + name: str, con, - schema=None, - if_exists="fail", - index=True, + schema: Optional[str] = None, + if_exists: str = "fail", + index: bool = True, index_label=None, - chunksize=None, + chunksize: Optional[int] = None, dtype=None, - method=None, + method: Optional[str] = None, ) -> None: """ Write records stored in a DataFrame to a SQL database. @@ -663,7 +669,7 @@ def to_sql( ) -def has_table(table_name, con, schema=None): +def has_table(table_name: str, con, schema: Optional[str] = None): """ Check if DataBase has named table. @@ -708,7 +714,9 @@ def _engine_builder(con): return con -def pandasSQL_builder(con, schema=None, meta=None, is_cursor=False): +def pandasSQL_builder( + con, schema: Optional[str] = None, meta=None, is_cursor: bool = False +): """ Convenience function to return the correct PandasSQL subclass based on the provided parameters. @@ -737,7 +745,7 @@ class SQLTable(PandasObject): def __init__( self, - name, + name: str, pandas_sql_engine, frame=None, index=True, @@ -795,7 +803,7 @@ def create(self): else: self._execute_create() - def _execute_insert(self, conn, keys, data_iter): + def _execute_insert(self, conn, keys: List[str], data_iter): """ Execute SQL statement inserting data @@ -810,7 +818,7 @@ def _execute_insert(self, conn, keys, data_iter): data = [dict(zip(keys, row)) for row in data_iter] conn.execute(self.table.insert(), data) - def _execute_insert_multi(self, conn, keys, data_iter): + def _execute_insert_multi(self, conn, keys: List[str], data_iter): """ Alternative to _execute_insert for DBs support multivalue INSERT. @@ -857,7 +865,7 @@ def insert_data(self): return column_names, data_list - def insert(self, chunksize=None, method=None): + def insert(self, chunksize: Optional[int] = None, method: Optional[str] = None): # set insert method if method is None: @@ -894,7 +902,12 @@ def insert(self, chunksize=None, method=None): exec_insert(conn, keys, chunk_iter) def _query_iterator( - self, result, chunksize, columns, coerce_float=True, parse_dates=None + self, + result, + chunksize: Optional[str], + columns, + coerce_float: bool = True, + parse_dates=None, ): """Return generator through chunked result set.""" while True: @@ -1203,7 +1216,7 @@ class SQLDatabase(PandasSQL): """ - def __init__(self, engine, schema=None, meta=None): + def __init__(self, engine, schema: Optional[str] = None, meta=None): self.connectable = engine if not meta: from sqlalchemy.schema import MetaData @@ -1228,13 +1241,13 @@ def execute(self, *args, **kwargs): def read_table( self, - table_name, - index_col=None, - coerce_float=True, + table_name: str, + index_col: Optional[Union[str, Sequence[str]]] = None, + coerce_float: bool = True, parse_dates=None, columns=None, - schema=None, - chunksize=None, + schema: Optional[str] = None, + chunksize: Optional[int] = None, ): """ Read SQL database table into a DataFrame. @@ -1288,7 +1301,12 @@ def read_table( @staticmethod def _query_iterator( - result, chunksize, columns, index_col=None, coerce_float=True, parse_dates=None + result, + chunksize: int, + columns, + index_col=None, + coerce_float=True, + parse_dates=None, ): """Return generator through chunked result set""" while True: @@ -1306,12 +1324,12 @@ def _query_iterator( def read_query( self, - sql, - index_col=None, - coerce_float=True, + sql: str, + index_col: Optional[str] = None, + coerce_float: bool = True, parse_dates=None, params=None, - chunksize=None, + chunksize: Optional[int] = None, ): """ Read SQL query into a DataFrame. @@ -1490,12 +1508,12 @@ def to_sql( def tables(self): return self.meta.tables - def has_table(self, name, schema=None): + def has_table(self, name: str, schema: Optional[str] = None): return self.connectable.run_callable( self.connectable.dialect.has_table, name, schema or self.meta.schema ) - def get_table(self, table_name, schema=None): + def get_table(self, table_name: str, schema: Optional[str] = None): schema = schema or self.meta.schema if schema: tbl = self.meta.tables.get(".".join([schema, table_name])) @@ -1511,7 +1529,7 @@ def get_table(self, table_name, schema=None): return tbl - def drop_table(self, table_name, schema=None): + def drop_table(self, table_name: str, schema: Optional[str] = None): schema = schema or self.meta.schema if self.has_table(table_name, schema): self.meta.reflect(only=[table_name], schema=schema) @@ -1608,7 +1626,7 @@ def _execute_create(self): for stmt in self.table: conn.execute(stmt) - def insert_statement(self, *, num_rows): + def insert_statement(self, *, num_rows: int): names = list(map(str, self.frame.columns)) wld = "?" # wildcard char escape = _get_valid_sqlite_name @@ -1737,7 +1755,7 @@ class SQLiteDatabase(PandasSQL): """ - def __init__(self, con, is_cursor=False): + def __init__(self, con, is_cursor: bool = False): self.is_cursor = is_cursor self.con = con @@ -1775,7 +1793,12 @@ def execute(self, *args, **kwargs): @staticmethod def _query_iterator( - cursor, chunksize, columns, index_col=None, coerce_float=True, parse_dates=None + cursor, + chunksize: int, + columns, + index_col=None, + coerce_float: bool = True, + parse_dates=None, ): """Return generator through chunked result set""" while True: @@ -1798,10 +1821,10 @@ def read_query( self, sql, index_col=None, - coerce_float=True, + coerce_float: bool = True, params=None, parse_dates=None, - chunksize=None, + chunksize: Optional[int] = None, ): args = _convert_params(sql, params) @@ -1908,7 +1931,7 @@ def to_sql( table.create() table.insert(chunksize, method) - def has_table(self, name, schema=None): + def has_table(self, name: str, schema: Optional[str] = None): # TODO(wesm): unused? # escape = _get_valid_sqlite_name # esc_name = escape(name) @@ -1918,14 +1941,21 @@ def has_table(self, name, schema=None): return len(self.execute(query, [name]).fetchall()) > 0 - def get_table(self, table_name, schema=None): + def get_table(self, table_name: str, schema: Optional[str] = None): return None # not supported in fallback mode - def drop_table(self, name, schema=None): + def drop_table(self, name: str, schema: Optional[str] = None): drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}" self.execute(drop_sql) - def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, schema=None): + def _create_sql_schema( + self, + frame, + table_name: str, + keys=None, + dtype=None, + schema: Optional[str] = None, + ): table = SQLiteTable( table_name, self, @@ -1938,7 +1968,9 @@ def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, schema=No return str(table.sql_schema()) -def get_schema(frame, name, keys=None, con=None, dtype=None, schema=None): +def get_schema( + frame, name: str, keys=None, con=None, dtype=None, schema: Optional[str] = None +): """ Get the SQL db table schema for the given frame. From 89f10fbf3c11c6824f714e57ddbf57c7b30a2951 Mon Sep 17 00:00:00 2001 From: Andrew Wieteska Date: Thu, 17 Dec 2020 19:00:16 -0500 Subject: [PATCH 2/3] review comment --- pandas/io/sql.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/io/sql.py b/pandas/io/sql.py index 66c4fbee3f1c9..89de9ab9426d2 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -129,7 +129,7 @@ def _wrap_result( columns, index_col=None, coerce_float: bool = True, - parse_dates: Optional[bool] = None, + parse_dates=None, ): """Wrap result set of query in a DataFrame.""" frame = DataFrame.from_records(data, columns=columns, coerce_float=coerce_float) From fbc518c37ef7425a1a0533054dbee02a85e0fe71 Mon Sep 17 00:00:00 2001 From: Andrew Wieteska Date: Sat, 19 Dec 2020 03:22:05 -0500 Subject: [PATCH 3/3] review comment --- pandas/io/sql.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pandas/io/sql.py b/pandas/io/sql.py index 89de9ab9426d2..836c16c109b57 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -7,7 +7,7 @@ from datetime import date, datetime, time from functools import partial import re -from typing import Iterator, List, Optional, Sequence, Union, overload +from typing import Any, Dict, Iterator, List, Optional, Sequence, Union, overload import warnings import numpy as np @@ -77,7 +77,9 @@ def _process_parse_dates_argument(parse_dates): return parse_dates -def _handle_date_column(col, utc: Optional[bool] = None, format: Optional[str] = None): +def _handle_date_column( + col, utc: Optional[bool] = None, format: Optional[Union[str, Dict[str, Any]]] = None +): if isinstance(format, dict): # GH35185 Allow custom error values in parse_dates argument of # read_sql like functions.