Add Literal typing for mode and projection_types (#2191)

Author: LeonLuttenberger

awswrangler/catalog/_create.py

@@ -1,7 +1,7 @@
 """AWS Glue Catalog Module."""
 
 import logging
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional
 
 import boto3
 
@@ -679,7 +679,7 @@ def create_parquet_table(
     description: Optional[str] = None,
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
-    mode: str = "overwrite",
+    mode: Literal["overwrite", "append"] = "overwrite",
     catalog_versioning: bool = False,
     transaction_id: Optional[str] = None,
     athena_partition_projection_settings: Optional[typing.AthenaPartitionProjectionSettings] = None,
@@ -840,7 +840,7 @@ def create_csv_table(  # pylint: disable=too-many-arguments,too-many-locals
     description: Optional[str] = None,
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
-    mode: str = "overwrite",
+    mode: Literal["overwrite", "append"] = "overwrite",
     catalog_versioning: bool = False,
     schema_evolution: bool = False,
     sep: str = ",",
@@ -1033,7 +1033,7 @@ def create_json_table(  # pylint: disable=too-many-arguments
     description: Optional[str] = None,
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
-    mode: str = "overwrite",
+    mode: Literal["overwrite", "append"] = "overwrite",
     catalog_versioning: bool = False,
     schema_evolution: bool = False,
     serde_library: Optional[str] = None,

awswrangler/data_quality/_create.py

@@ -3,7 +3,7 @@
 import logging
 import pprint
 import uuid
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import boto3
 
@@ -121,7 +121,7 @@ def create_ruleset(
 @apply_configs
 def update_ruleset(
     name: str,
-    mode: str = "overwrite",
+    mode: Literal["overwrite", "upsert"] = "overwrite",
     df_rules: Optional[pd.DataFrame] = None,
     dqdl_rules: Optional[str] = None,
     description: str = "",

awswrangler/emr.py

@@ -3,7 +3,7 @@
 
 import logging
 import pprint
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import boto3
 
@@ -12,6 +12,9 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+_ActionOnFailureLiteral = Literal["TERMINATE_JOB_FLOW", "TERMINATE_CLUSTER", "CANCEL_AND_WAIT", "CONTINUE"]
+
+
 def _get_ecr_credentials_refresh_content(region: str) -> str:
     return f"""
 import subprocess
@@ -815,7 +818,7 @@ def submit_step(
     cluster_id: str,
     command: str,
     name: str = "my-step",
-    action_on_failure: str = "CONTINUE",
+    action_on_failure: _ActionOnFailureLiteral = "CONTINUE",
     script: bool = False,
     boto3_session: Optional[boto3.Session] = None,
 ) -> str:
@@ -865,7 +868,7 @@ def submit_step(
 def build_step(
     command: str,
     name: str = "my-step",
-    action_on_failure: str = "CONTINUE",
+    action_on_failure: _ActionOnFailureLiteral = "CONTINUE",
     script: bool = False,
     region: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
@@ -951,7 +954,10 @@ def get_step_state(cluster_id: str, step_id: str, boto3_session: Optional[boto3.
 
 
 def submit_ecr_credentials_refresh(
-    cluster_id: str, path: str, action_on_failure: str = "CONTINUE", boto3_session: Optional[boto3.Session] = None
+    cluster_id: str,
+    path: str,
+    action_on_failure: _ActionOnFailureLiteral = "CONTINUE",
+    boto3_session: Optional[boto3.Session] = None,
 ) -> str:
     """Update internal ECR credentials.
 
@@ -999,10 +1005,10 @@ def submit_ecr_credentials_refresh(
 def build_spark_step(
     path: str,
     args: Optional[List[str]] = None,
-    deploy_mode: str = "cluster",
+    deploy_mode: Literal["cluster", "client"] = "cluster",
     docker_image: Optional[str] = None,
     name: str = "my-step",
-    action_on_failure: str = "CONTINUE",
+    action_on_failure: _ActionOnFailureLiteral = "CONTINUE",
     region: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> Dict[str, Any]:
@@ -1074,10 +1080,10 @@ def submit_spark_step(
     cluster_id: str,
     path: str,
     args: Optional[List[str]] = None,
-    deploy_mode: str = "cluster",
+    deploy_mode: Literal["cluster", "client"] = "cluster",
     docker_image: Optional[str] = None,
     name: str = "my-step",
-    action_on_failure: str = "CONTINUE",
+    action_on_failure: _ActionOnFailureLiteral = "CONTINUE",
     region: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> str:

awswrangler/mysql.py

@@ -3,7 +3,7 @@
 
 import logging
 import uuid
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, overload
+from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Type, Union, cast, overload
 
 import boto3
 import pyarrow as pa
@@ -402,14 +402,19 @@ def read_sql_table(
     )
 
 
+_ToSqlModeLiteral = Literal[
+    "append", "overwrite", "upsert_replace_into", "upsert_duplicate_key", "upsert_distinct", "ignore"
+]
+
+
 @_utils.check_optional_dependency(pymysql, "pymysql")
 @apply_configs
 def to_sql(
     df: pd.DataFrame,
     con: "pymysql.connections.Connection[Any]",
     table: str,
     schema: str,
-    mode: str = "append",
+    mode: _ToSqlModeLiteral = "append",
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
     varchar_lengths: Optional[Dict[str, int]] = None,
@@ -430,7 +435,7 @@ def to_sql(
     schema : str
         Schema name
     mode : str
-        Append, overwrite, upsert_duplicate_key, upsert_replace_into, upsert_distinct, ignore.
+        append, overwrite, upsert_duplicate_key, upsert_replace_into, upsert_distinct, ignore.
             append: Inserts new records into table.
             overwrite: Drops table and recreates.
             upsert_duplicate_key: Performs an upsert using `ON DUPLICATE KEY` clause. Requires table schema to have
@@ -484,7 +489,7 @@ def to_sql(
     if df.empty is True:
         raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
 
-    mode = mode.strip().lower()
+    mode = cast(_ToSqlModeLiteral, mode.strip().lower())
     allowed_modes = [
         "append",
         "overwrite",

awswrangler/oracle.py

@@ -3,7 +3,19 @@
 
 import logging
 from decimal import Decimal
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar, Union, overload
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    overload,
+)
 
 import boto3
 import pyarrow as pa
@@ -404,7 +416,7 @@ def to_sql(
     con: "oracledb.Connection",
     table: str,
     schema: str,
-    mode: str = "append",
+    mode: Literal["append", "overwrite"] = "append",
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
     varchar_lengths: Optional[Dict[str, int]] = None,

awswrangler/postgresql.py

@@ -3,7 +3,7 @@
 
 import logging
 from ssl import SSLContext
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, overload
+from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, cast, overload
 
 import boto3
 import pyarrow as pa
@@ -393,14 +393,17 @@ def read_sql_table(
     )
 
 
+_ToSqlModeLiteral = Literal["append", "overwrite", "upsert"]
+
+
 @_utils.check_optional_dependency(pg8000, "pg8000")
 @apply_configs
 def to_sql(
     df: pd.DataFrame,
     con: "pg8000.Connection",
     table: str,
     schema: str,
-    mode: str = "append",
+    mode: _ToSqlModeLiteral = "append",
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
     varchar_lengths: Optional[Dict[str, int]] = None,
@@ -472,7 +475,7 @@ def to_sql(
     if df.empty is True:
         raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
 
-    mode = mode.strip().lower()
+    mode = cast(_ToSqlModeLiteral, mode.strip().lower())
     allowed_modes = ["append", "overwrite", "upsert"]
     _db_utils.validate_mode(mode=mode, allowed_modes=allowed_modes)
     if mode == "upsert" and not upsert_conflict_columns:

awswrangler/quicksight/_create.py

@@ -2,7 +2,7 @@
 
 import logging
 import uuid
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import boto3
 
@@ -202,7 +202,7 @@ def create_athena_dataset(
     sql_name: Optional[str] = None,
     data_source_name: Optional[str] = None,
     data_source_arn: Optional[str] = None,
-    import_mode: str = "DIRECT_QUERY",
+    import_mode: Literal["SPICE", "DIRECT_QUERY"] = "DIRECT_QUERY",
     allowed_to_use: Optional[List[str]] = None,
     allowed_to_manage: Optional[List[str]] = None,
     logical_table_alias: str = "LogicalTable",

awswrangler/redshift.py

@@ -5,7 +5,7 @@
 import json
 import logging
 import uuid
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, overload
+from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, overload
 
 import boto3
 import botocore
@@ -24,6 +24,11 @@
 _RS_DISTSTYLES: List[str] = ["AUTO", "EVEN", "ALL", "KEY"]
 _RS_SORTSTYLES: List[str] = ["COMPOUND", "INTERLEAVED"]
 
+_ToSqlModeLiteral = Literal["append", "overwrite", "upsert"]
+_ToSqlOverwriteModeLiteral = Literal["drop", "cascade", "truncate", "delete"]
+_ToSqlDistStyleLiteral = Literal["AUTO", "EVEN", "ALL", "KEY"]
+_ToSqlSortStyleLiteral = Literal["COMPOUND", "INTERLEAVED"]
+
 
 def _validate_connection(con: "redshift_connector.Connection") -> None:
     if not isinstance(con, redshift_connector.Connection):
@@ -909,13 +914,13 @@ def to_sql(  # pylint: disable=too-many-locals
     con: "redshift_connector.Connection",
     table: str,
     schema: str,
-    mode: str = "append",
-    overwrite_method: str = "drop",
+    mode: _ToSqlModeLiteral = "append",
+    overwrite_method: _ToSqlOverwriteModeLiteral = "drop",
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
-    diststyle: str = "AUTO",
+    diststyle: _ToSqlDistStyleLiteral = "AUTO",
     distkey: Optional[str] = None,
-    sortstyle: str = "COMPOUND",
+    sortstyle: _ToSqlSortStyleLiteral = "COMPOUND",
     sortkey: Optional[List[str]] = None,
     primary_keys: Optional[List[str]] = None,
     varchar_lengths_default: int = 256,
@@ -1084,7 +1089,7 @@ def unload_to_files(
     aws_secret_access_key: Optional[str] = None,
     aws_session_token: Optional[str] = None,
     region: Optional[str] = None,
-    unload_format: Optional[str] = None,
+    unload_format: Optional[Literal["CSV", "PARQUET"]] = None,
     max_file_size: Optional[float] = None,
     kms_key_id: Optional[str] = None,
     manifest: bool = False,
@@ -1380,11 +1385,11 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
     aws_secret_access_key: Optional[str] = None,
     aws_session_token: Optional[str] = None,
     parquet_infer_sampling: float = 1.0,
-    mode: str = "append",
-    overwrite_method: str = "drop",
-    diststyle: str = "AUTO",
+    mode: _ToSqlModeLiteral = "append",
+    overwrite_method: _ToSqlOverwriteModeLiteral = "drop",
+    diststyle: _ToSqlDistStyleLiteral = "AUTO",
     distkey: Optional[str] = None,
-    sortstyle: str = "COMPOUND",
+    sortstyle: _ToSqlSortStyleLiteral = "COMPOUND",
     sortkey: Optional[List[str]] = None,
     primary_keys: Optional[List[str]] = None,
     varchar_lengths_default: int = 256,
@@ -1608,11 +1613,11 @@ def copy(  # pylint: disable=too-many-arguments,too-many-locals
     aws_session_token: Optional[str] = None,
     index: bool = False,
     dtype: Optional[Dict[str, str]] = None,
-    mode: str = "append",
-    overwrite_method: str = "drop",
-    diststyle: str = "AUTO",
+    mode: _ToSqlModeLiteral = "append",
+    overwrite_method: _ToSqlOverwriteModeLiteral = "drop",
+    diststyle: _ToSqlDistStyleLiteral = "AUTO",
     distkey: Optional[str] = None,
-    sortstyle: str = "COMPOUND",
+    sortstyle: _ToSqlSortStyleLiteral = "COMPOUND",
     sortkey: Optional[List[str]] = None,
     primary_keys: Optional[List[str]] = None,
     varchar_lengths_default: int = 256,

awswrangler/s3/_copy.py

@@ -2,7 +2,7 @@
 
 import itertools
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
 
 import boto3
 from boto3.s3.transfer import TransferConfig
@@ -73,7 +73,7 @@ def _copy(
 def merge_datasets(
     source_path: str,
     target_path: str,
-    mode: str = "append",
+    mode: Literal["append", "overwrite", "overwrite_partitions"] = "append",
     ignore_empty: bool = False,
     use_threads: Union[bool, int] = True,
     boto3_session: Optional[boto3.Session] = None,

awswrangler/s3/_write_parquet.py

@@ -4,7 +4,7 @@
 import math
 import uuid
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, cast
 
 import boto3
 import pandas as pd
@@ -241,7 +241,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-b
     partition_cols: Optional[List[str]] = None,
     bucketing_info: Optional[BucketingInfoTuple] = None,
     concurrent_partitioning: bool = False,
-    mode: Optional[str] = None,
+    mode: Optional[Literal["append", "overwrite", "overwrite_partitions"]] = None,
     catalog_versioning: bool = False,
     schema_evolution: bool = True,
     database: Optional[str] = None,
@@ -821,7 +821,7 @@ def store_parquet_metadata(  # pylint: disable=too-many-arguments,too-many-local
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
     compression: Optional[str] = None,
-    mode: str = "overwrite",
+    mode: Literal["append", "overwrite"] = "overwrite",
     catalog_versioning: bool = False,
     regular_partitions: bool = True,
     athena_partition_projection_settings: Optional[typing.AthenaPartitionProjectionSettings] = None,