Applying mypy strict mode.

Author: igorborgest

awswrangler/__init__.py

@@ -11,4 +11,21 @@
 from awswrangler.__metadata__ import __description__, __license__, __title__, __version__  # noqa
 from awswrangler._config import config  # noqa
 
+__all__ = [
+    "athena",
+    "catalog",
+    "cloudwatch",
+    "db",
+    "emr",
+    "exceptions",
+    "quicksight",
+    "s3",
+    "sts",
+    "config",
+    "__description__",
+    "__license__",
+    "__title__",
+    "__version__",
+]
+
 _logging.getLogger("awswrangler").addHandler(_logging.NullHandler())

awswrangler/_config.py

@@ -16,7 +16,7 @@
 
 
 class _ConfigArg(NamedTuple):
-    dtype: Type
+    dtype: Type[Union[str, bool, int]]
     nullable: bool
     enforced: bool = False
 
@@ -36,7 +36,7 @@ class _ConfigArg(NamedTuple):
 class _Config:
     """Wrangler's Configuration class."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._loaded_values: Dict[str, _ConfigValueType] = {}
         name: str
         for name in _CONFIG_ARGS:
@@ -127,11 +127,11 @@ def _reset_item(self, item: str) -> None:
             del self._loaded_values[item]
         self._load_config(name=item)
 
-    def _repr_html_(self):
+    def _repr_html_(self) -> Any:
         return self.to_pandas().to_html()
 
     @staticmethod
-    def _apply_type(name: str, value: Any, dtype: Type, nullable: bool) -> _ConfigValueType:
+    def _apply_type(name: str, value: Any, dtype: Type[Union[str, bool, int]], nullable: bool) -> _ConfigValueType:
         if _Config._is_null(value=value):
             if nullable is True:
                 return None
@@ -215,7 +215,9 @@ def s3fs_block_size(self, value: int) -> None:
         self._set_config_value(key="s3fs_block_size", value=value)
 
 
-def _inject_config_doc(doc: str, available_configs: Tuple[str, ...]) -> str:
+def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -> str:
+    if doc is None:
+        return "Undocumented function."
     if "\n    Parameters" not in doc:
         return doc
     header: str = (
@@ -235,14 +237,14 @@ def _inject_config_doc(doc: str, available_configs: Tuple[str, ...]) -> str:
     return _utils.insert_str(text=doc, token="\n    Parameters", insert=insertion)
 
 
-def apply_configs(function) -> Callable:
+def apply_configs(function: Callable[..., Any]) -> Callable[..., Any]:
     """Decorate some function with configs."""
     signature = inspect.signature(function)
     args_names: Tuple[str, ...] = tuple(signature.parameters.keys())
     available_configs: Tuple[str, ...] = tuple(x for x in _CONFIG_ARGS if x in args_names)
 
-    def wrapper(*args, **kwargs):
-        args: Dict[str, Any] = signature.bind_partial(*args, **kwargs).arguments
+    def wrapper(*args_raw: Any, **kwargs: Any) -> Any:
+        args: Dict[str, Any] = signature.bind_partial(*args_raw, **kwargs).arguments
         for name in available_configs:
             if hasattr(config, name) is True:
                 value: _ConfigValueType = config[name]

awswrangler/_data_types.py

@@ -315,7 +315,7 @@ def process_not_inferred_dtype(ex: pa.ArrowInvalid) -> pa.DataType:
     """Infer data type from PyArrow inference exception."""
     ex_str = str(ex)
     _logger.debug("PyArrow was not able to infer data type:\n%s", ex_str)
-    match: Optional[Match] = re.search(
+    match: Optional[Match[str]] = re.search(
         pattern="Could not convert (.*) with type (.*): did not recognize "
         "Python value type when inferring an Arrow data type",
         string=ex_str,

awswrangler/_utils.py

@@ -6,7 +6,7 @@
 import os
 import random
 import time
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
 
 import boto3  # type: ignore
 import botocore.config  # type: ignore
@@ -196,19 +196,19 @@ def get_fs(
     return fs
 
 
-def open_file(fs: s3fs.S3FileSystem, **kwargs) -> Any:
+def open_file(fs: s3fs.S3FileSystem, **kwargs: Any) -> Any:
     """Open s3fs file with retries to overcome eventual consistency."""
     fs.invalidate_cache()
     fs.clear_instance_cache()
     return try_it(f=fs.open, ex=FileNotFoundError, **kwargs)
 
 
-def empty_generator() -> Generator:
+def empty_generator() -> Generator[None, None, None]:
     """Empty Generator."""
     yield from ()
 
 
-def ensure_postgresql_casts():
+def ensure_postgresql_casts() -> None:
     """Ensure that psycopg2 will handle some data types right."""
     psycopg2.extensions.register_adapter(bytes, psycopg2.Binary)
     typecast_bytea = lambda data, cur: None if data is None else bytes(psycopg2.BINARY(data, cur))  # noqa
@@ -225,7 +225,7 @@ def get_region_from_subnet(subnet_id: str, boto3_session: Optional[boto3.Session
     """Extract region from Subnet ID."""
     session: boto3.Session = ensure_session(session=boto3_session)
     client_ec2: boto3.client = client(service_name="ec2", session=session)
-    return client_ec2.describe_subnets(SubnetIds=[subnet_id])["Subnets"][0]["AvailabilityZone"][:-1]
+    return cast(str, client_ec2.describe_subnets(SubnetIds=[subnet_id])["Subnets"][0]["AvailabilityZone"][:-1])
 
 
 def get_region_from_session(boto3_session: Optional[boto3.Session] = None, default_region: Optional[str] = None) -> str:
@@ -279,7 +279,7 @@ def check_duplicated_columns(df: pd.DataFrame) -> Any:
         raise exceptions.InvalidDataFrame(f"There is duplicated column names in your DataFrame: {duplicated}")
 
 
-def try_it(f: Callable, ex, base: float = 1.0, max_num_tries: int = 3, **kwargs) -> Any:
+def try_it(f: Callable[..., Any], ex: Any, base: float = 1.0, max_num_tries: int = 3, **kwargs: Any) -> Any:
     """Run function with decorrelated Jitter.
 
     Reference: https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/

awswrangler/athena/__init__.py

@@ -13,3 +13,18 @@
     stop_query_execution,
     wait_query,
 )
+
+__all__ = [
+    "read_sql_query",
+    "read_sql_table",
+    "create_athena_bucket",
+    "describe_table",
+    "get_query_columns_types",
+    "get_query_execution",
+    "get_work_group",
+    "repair_table",
+    "show_create_table",
+    "start_query_execution",
+    "stop_query_execution",
+    "wait_query",
+]

awswrangler/athena/_read.py

@@ -111,7 +111,7 @@ def _get_last_query_executions(
         yield execution_data.get("QueryExecutions")
 
 
-def _sort_successful_executions_data(query_executions: List[Dict[str, Any]]):
+def _sort_successful_executions_data(query_executions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """
     Sorts `_get_last_query_executions`'s results based on query Completion DateTime.
 
@@ -289,14 +289,16 @@ def _fetch_csv_result(
 
 
 def _resolve_query_with_cache(
-    cache_info,
+    cache_info: _CacheInfo,
     categories: Optional[List[str]],
     chunksize: Optional[Union[int, bool]],
     use_threads: bool,
     session: Optional[boto3.Session],
-):
+) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Fetch cached data and return it as a pandas DataFrame (or list of DataFrames)."""
     _logger.debug("cache_info:\n%s", cache_info)
+    if cache_info.query_execution_id is None:
+        raise RuntimeError("Trying to resolve with cache but w/o any query execution ID.")
     query_metadata: _QueryMetadata = _get_query_metadata(
         query_execution_id=cache_info.query_execution_id,
         boto3_session=session,
@@ -330,14 +332,14 @@ def _resolve_query_without_cache_ctas(
     keep_files: bool,
     chunksize: Union[int, bool, None],
     categories: Optional[List[str]],
-    encryption,
+    encryption: Optional[str],
     workgroup: Optional[str],
     kms_key: Optional[str],
     wg_config: _WorkGroupConfig,
     name: Optional[str],
     use_threads: bool,
     boto3_session: boto3.Session,
-):
+) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     path: str = f"{s3_output}/{name}"
     ext_location: str = "\n" if wg_config.enforced is True else f",\n    external_location = '{path}'\n"
     sql = (
@@ -406,13 +408,13 @@ def _resolve_query_without_cache_regular(
     keep_files: bool,
     chunksize: Union[int, bool, None],
     categories: Optional[List[str]],
-    encryption,
+    encryption: Optional[str],
     workgroup: Optional[str],
     kms_key: Optional[str],
     wg_config: _WorkGroupConfig,
     use_threads: bool,
     boto3_session: boto3.Session,
-):
+) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     _logger.debug("sql: %s", sql)
     query_id: str = _start_query_execution(
         sql=sql,
@@ -698,7 +700,7 @@ def read_sql_table(
     table: str,
     database: str,
     ctas_approach: bool = True,
-    categories: List[str] = None,
+    categories: Optional[List[str]] = None,
     chunksize: Optional[Union[int, bool]] = None,
     s3_output: Optional[str] = None,
     workgroup: Optional[str] = None,

awswrangler/athena/_utils.py

@@ -5,7 +5,7 @@
 import time
 import warnings
 from decimal import Decimal
-from typing import Any, Dict, List, NamedTuple, Optional, Union
+from typing import Any, Dict, Generator, List, NamedTuple, Optional, Union, cast
 
 import boto3  # type: ignore
 import pandas as pd  # type: ignore
@@ -89,7 +89,7 @@ def _start_query_execution(
     client_athena: boto3.client = _utils.client(service_name="athena", session=session)
     _logger.debug("args: \n%s", pprint.pformat(args))
     response: Dict[str, Any] = client_athena.start_query_execution(**args)
-    return response["QueryExecutionId"]
+    return cast(str, response["QueryExecutionId"])
 
 
 def _get_workgroup_config(session: boto3.Session, workgroup: Optional[str] = None) -> _WorkGroupConfig:
@@ -137,7 +137,7 @@ def _fetch_txt_result(query_metadata: _QueryMetadata, keep_files: bool, boto3_se
 
 def _parse_describe_table(df: pd.DataFrame) -> pd.DataFrame:
     origin_df_dict = df.to_dict()
-    target_df_dict: Dict[str, List] = {"Column Name": [], "Type": [], "Partition": [], "Comment": []}
+    target_df_dict: Dict[str, List[Union[str, bool]]] = {"Column Name": [], "Type": [], "Partition": [], "Comment": []}
     for index, col_name in origin_df_dict["col_name"].items():
         col_name = col_name.strip()
         if col_name.startswith("#") or col_name == "":
@@ -156,7 +156,7 @@ def _parse_describe_table(df: pd.DataFrame) -> pd.DataFrame:
 def _get_query_metadata(  # pylint: disable=too-many-statements
     query_execution_id: str,
     boto3_session: boto3.Session,
-    categories: List[str] = None,
+    categories: Optional[List[str]] = None,
     query_execution_payload: Optional[Dict[str, Any]] = None,
 ) -> _QueryMetadata:
     """Get query metadata."""
@@ -226,7 +226,9 @@ def _get_query_metadata(  # pylint: disable=too-many-statements
     return query_metadata
 
 
-def _empty_dataframe_response(chunked: bool, query_metadata: _QueryMetadata):
+def _empty_dataframe_response(
+    chunked: bool, query_metadata: _QueryMetadata
+) -> Union[pd.DataFrame, Generator[None, None, None]]:
     """Generate an empty dataframe response."""
     if chunked is False:
         df = pd.DataFrame()
@@ -425,7 +427,7 @@ def repair_table(
         boto3_session=session,
     )
     response: Dict[str, Any] = wait_query(query_execution_id=query_id, boto3_session=session)
-    return response["Status"]["State"]
+    return cast(str, response["Status"]["State"])
 
 
 @apply_configs
@@ -556,7 +558,7 @@ def show_create_table(
     )
     query_metadata: _QueryMetadata = _get_query_metadata(query_execution_id=query_id, boto3_session=session)
     raw_result = _fetch_txt_result(query_metadata=query_metadata, keep_files=True, boto3_session=session,)
-    return raw_result.createtab_stmt.str.strip().str.cat(sep=" ")
+    return cast(str, raw_result.createtab_stmt.str.strip().str.cat(sep=" "))
 
 
 def get_work_group(workgroup: str, boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]:
@@ -581,7 +583,7 @@ def get_work_group(workgroup: str, boto3_session: Optional[boto3.Session] = None
 
     """
     client_athena: boto3.client = _utils.client(service_name="athena", session=boto3_session)
-    return client_athena.get_work_group(WorkGroup=workgroup)
+    return cast(Dict[str, Any], client_athena.get_work_group(WorkGroup=workgroup))
 
 
 def stop_query_execution(query_execution_id: str, boto3_session: Optional[boto3.Session] = None) -> None:
@@ -645,7 +647,7 @@ def wait_query(query_execution_id: str, boto3_session: Optional[boto3.Session] =
         raise exceptions.QueryFailed(response["QueryExecution"]["Status"].get("StateChangeReason"))
     if state == "CANCELLED":
         raise exceptions.QueryCancelled(response["QueryExecution"]["Status"].get("StateChangeReason"))
-    return response["QueryExecution"]
+    return cast(Dict[str, Any], response["QueryExecution"])
 
 
 def get_query_execution(query_execution_id: str, boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]:
@@ -673,4 +675,4 @@ def get_query_execution(query_execution_id: str, boto3_session: Optional[boto3.S
     """
     client_athena: boto3.client = _utils.client(service_name="athena", session=boto3_session)
     response: Dict[str, Any] = client_athena.get_query_execution(QueryExecutionId=query_execution_id)
-    return response["QueryExecution"]
+    return cast(Dict[str, Any], response["QueryExecution"])

awswrangler/catalog/__init__.py

@@ -40,3 +40,42 @@
     sanitize_dataframe_columns_names,
     sanitize_table_name,
 )
+
+__all__ = [
+    "add_csv_partitions",
+    "add_parquet_partitions",
+    "does_table_exist",
+    "drop_duplicated_columns",
+    "extract_athena_types",
+    "sanitize_column_name",
+    "sanitize_dataframe_columns_names",
+    "sanitize_table_name",
+    "_create_csv_table",
+    "_create_parquet_table",
+    "create_csv_table",
+    "create_database",
+    "create_parquet_table",
+    "overwrite_table_parameters",
+    "upsert_table_parameters",
+    "_get_table_input",
+    "databases",
+    "get_columns_comments",
+    "get_connection",
+    "get_csv_partitions",
+    "get_databases",
+    "get_engine",
+    "get_parquet_partitions",
+    "get_partitions",
+    "get_table_description",
+    "get_table_location",
+    "get_table_number_of_versions",
+    "get_table_parameters",
+    "get_table_types",
+    "get_table_versions",
+    "get_tables",
+    "search_tables",
+    "table",
+    "tables",
+    "delete_database",
+    "delete_table_if_exists",
+]

awswrangler/catalog/_add.py

@@ -19,7 +19,7 @@ def _add_partitions(
     boto3_session: Optional[boto3.Session],
     inputs: List[Dict[str, Any]],
     catalog_id: Optional[str] = None,
-):
+) -> None:
     chunks: List[List[Dict[str, Any]]] = _utils.chunkify(lst=inputs, max_length=100)
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
     for chunk in chunks:  # pylint: disable=too-many-nested-blocks

awswrangler/catalog/_create.py

@@ -43,7 +43,7 @@ def _create_table(  # pylint: disable=too-many-branches,too-many-statements
     projection_intervals: Optional[Dict[str, str]],
     projection_digits: Optional[Dict[str, str]],
     catalog_id: Optional[str],
-):
+) -> None:
     # Description
     mode = _update_if_necessary(dic=table_input, key="Description", value=description, mode=mode)