TYP: pandas/io/sql.py (easy: bool/str) (pandas-dev#38537)

Author: arw2019

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 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=None, format=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.
@@ -124,7 +126,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=None,
+):
     """Wrap result set of query in a DataFrame."""
     frame = DataFrame.from_records(data, columns=columns, coerce_float=coerce_float)
 
@@ -197,11 +205,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 +329,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 +428,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 +590,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 +671,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 +716,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 +747,7 @@ class SQLTable(PandasObject):
 
     def __init__(
         self,
-        name,
+        name: str,
         pandas_sql_engine,
         frame=None,
         index=True,
@@ -795,7 +805,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 +820,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 +867,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 +904,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 +1218,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 +1243,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 +1303,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 +1326,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 +1510,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 +1531,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 +1628,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 +1757,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 +1795,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 +1823,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 +1933,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 +1943,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 +1970,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.