Add optional typecheck for nullable parameters (#4767)

jerrypeng7773 · web-flow · commit 7c9f870054fa · 2024-07-31T16:01:28.000-07:00
* Add optional typecheck for nullable parameters

* Add optional typecheck for nullable parameters
diff --git a/src/sagemaker/workflow/automl_step.py b/src/sagemaker/workflow/automl_step.py
@@ -34,11 +34,11 @@ def __init__(
         self,
         name: str,
         step_args: _JobStepArguments,
-        display_name: str = None,
-        description: str = None,
-        cache_config: CacheConfig = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
-        retry_policies: List[RetryPolicy] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
     ):
         """Construct a `AutoMLStep`, given a `AutoML` instance.
 
diff --git a/src/sagemaker/workflow/callback_step.py b/src/sagemaker/workflow/callback_step.py
@@ -84,9 +84,9 @@ def __init__(
         sqs_queue_url: str,
         inputs: dict,
         outputs: List[CallbackOutput],
-        display_name: str = None,
-        description: str = None,
-        cache_config: CacheConfig = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a CallbackStep.
diff --git a/src/sagemaker/workflow/clarify_check_step.py b/src/sagemaker/workflow/clarify_check_step.py
@@ -159,11 +159,11 @@ def __init__(
         skip_check: Union[bool, PipelineVariable] = False,
         fail_on_violation: Union[bool, PipelineVariable] = True,
         register_new_baseline: Union[bool, PipelineVariable] = False,
-        model_package_group_name: Union[str, PipelineVariable] = None,
-        supplied_baseline_constraints: Union[str, PipelineVariable] = None,
-        display_name: str = None,
-        description: str = None,
-        cache_config: CacheConfig = None,
+        model_package_group_name: Optional[Union[str, PipelineVariable]] = None,
+        supplied_baseline_constraints: Optional[Union[str, PipelineVariable]] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a ClarifyCheckStep.
diff --git a/src/sagemaker/workflow/condition_step.py b/src/sagemaker/workflow/condition_step.py
@@ -20,7 +20,6 @@
 
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.functions import JsonGet as NewJsonGet
-from sagemaker.workflow.step_outputs import StepOutput
 from sagemaker.workflow.steps import (
     Step,
     StepTypeEnum,
@@ -41,11 +40,11 @@ def __init__(
         self,
         name: str,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
-        display_name: str = None,
-        description: str = None,
-        conditions: List[Condition] = None,
-        if_steps: List[Union[Step, StepCollection, StepOutput]] = None,
-        else_steps: List[Union[Step, StepCollection, StepOutput]] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        conditions: Optional[List[Condition]] = None,
+        if_steps: Optional[List[Union[Step, StepCollection]]] = None,
+        else_steps: Optional[List[Union[Step, StepCollection]]] = None,
     ):
         """Construct a ConditionStep for pipelines to support conditional branching.
 
diff --git a/src/sagemaker/workflow/emr_step.py b/src/sagemaker/workflow/emr_step.py
@@ -161,9 +161,9 @@ def __init__(
         cluster_id: str,
         step_config: EMRStepConfig,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
-        cache_config: CacheConfig = None,
-        cluster_config: Dict[str, Any] = None,
-        execution_role_arn: str = None,
+        cache_config: Optional[CacheConfig] = None,
+        cluster_config: Optional[Dict[str, Any]] = None,
+        execution_role_arn: Optional[str] = None,
     ):
         """Constructs an `EMRStep`.
 
diff --git a/src/sagemaker/workflow/fail_step.py b/src/sagemaker/workflow/fail_step.py
@@ -29,9 +29,9 @@ class FailStep(Step):
     def __init__(
         self,
         name: str,
-        error_message: Union[str, PipelineVariable] = None,
-        display_name: str = None,
-        description: str = None,
+        error_message: Optional[Union[str, PipelineVariable]] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a `FailStep`.
diff --git a/src/sagemaker/workflow/lambda_step.py b/src/sagemaker/workflow/lambda_step.py
@@ -84,11 +84,11 @@ def __init__(
         self,
         name: str,
         lambda_func: Lambda,
-        display_name: str = None,
-        description: str = None,
-        inputs: dict = None,
-        outputs: List[LambdaOutput] = None,
-        cache_config: CacheConfig = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        inputs: Optional[dict] = None,
+        outputs: Optional[List[LambdaOutput]] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a LambdaStep.
diff --git a/src/sagemaker/workflow/monitor_batch_transform_step.py b/src/sagemaker/workflow/monitor_batch_transform_step.py
@@ -48,8 +48,8 @@ def __init__(
         check_job_configuration: CheckJobConfig,
         monitor_before_transform: bool = False,
         fail_on_violation: Union[bool, PipelineVariable] = True,
-        supplied_baseline_statistics: Union[str, PipelineVariable] = None,
-        supplied_baseline_constraints: Union[str, PipelineVariable] = None,
+        supplied_baseline_statistics: Optional[Union[str, PipelineVariable]] = None,
+        supplied_baseline_constraints: Optional[Union[str, PipelineVariable]] = None,
         display_name: Optional[str] = None,
         description: Optional[str] = None,
     ):
diff --git a/src/sagemaker/workflow/quality_check_step.py b/src/sagemaker/workflow/quality_check_step.py
@@ -125,12 +125,12 @@ def __init__(
         skip_check: Union[bool, PipelineVariable] = False,
         fail_on_violation: Union[bool, PipelineVariable] = True,
         register_new_baseline: Union[bool, PipelineVariable] = False,
-        model_package_group_name: Union[str, PipelineVariable] = None,
-        supplied_baseline_statistics: Union[str, PipelineVariable] = None,
-        supplied_baseline_constraints: Union[str, PipelineVariable] = None,
-        display_name: str = None,
-        description: str = None,
-        cache_config: CacheConfig = None,
+        model_package_group_name: Optional[Union[str, PipelineVariable]] = None,
+        supplied_baseline_statistics: Optional[Union[str, PipelineVariable]] = None,
+        supplied_baseline_constraints: Optional[Union[str, PipelineVariable]] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a QualityCheckStep.
diff --git a/src/sagemaker/workflow/retry.py b/src/sagemaker/workflow/retry.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from enum import Enum
-from typing import List
+from typing import List, Optional
 import attr
 
 from sagemaker.workflow.entities import Entity, DefaultEnumMeta, RequestType
@@ -133,8 +133,8 @@ def __init__(
         exception_types: List[StepExceptionTypeEnum],
         backoff_rate: float = 2.0,
         interval_seconds: int = 1,
-        max_attempts: int = None,
-        expire_after_mins: int = None,
+        max_attempts: Optional[int] = None,
+        expire_after_mins: Optional[int] = None,
     ):
         super().__init__(backoff_rate, interval_seconds, max_attempts, expire_after_mins)
         for exception_type in exception_types:
@@ -177,12 +177,12 @@ class SageMakerJobStepRetryPolicy(RetryPolicy):
 
     def __init__(
         self,
-        exception_types: List[SageMakerJobExceptionTypeEnum] = None,
-        failure_reason_types: List[SageMakerJobExceptionTypeEnum] = None,
+        exception_types: Optional[List[SageMakerJobExceptionTypeEnum]] = None,
+        failure_reason_types: Optional[List[SageMakerJobExceptionTypeEnum]] = None,
         backoff_rate: float = 2.0,
         interval_seconds: int = 1,
-        max_attempts: int = None,
-        expire_after_mins: int = None,
+        max_attempts: Optional[int] = None,
+        expire_after_mins: Optional[int] = None,
     ):
         super().__init__(backoff_rate, interval_seconds, max_attempts, expire_after_mins)
 
diff --git a/src/sagemaker/workflow/selective_execution_config.py b/src/sagemaker/workflow/selective_execution_config.py
@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """Pipeline Parallelism Configuration"""
 from __future__ import absolute_import
-from typing import List
+from typing import List, Optional
 from sagemaker.workflow.entities import RequestType
 
 
@@ -25,8 +25,8 @@ class SelectiveExecutionConfig:
     def __init__(
         self,
         selected_steps: List[str],
-        source_pipeline_execution_arn: str = None,
         reference_latest_execution: bool = True,
+        source_pipeline_execution_arn: Optional[str] = None,
     ):
         """Create a `SelectiveExecutionConfig`.
 
diff --git a/src/sagemaker/workflow/step_collections.py b/src/sagemaker/workflow/step_collections.py
@@ -72,11 +72,11 @@ def __init__(
         response_types,
         inference_instances=None,
         transform_instances=None,
-        estimator: EstimatorBase = None,
+        estimator: Optional[EstimatorBase] = None,
         model_data=None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
-        repack_model_step_retry_policies: List[RetryPolicy] = None,
-        register_model_step_retry_policies: List[RetryPolicy] = None,
+        repack_model_step_retry_policies: Optional[List[RetryPolicy]] = None,
+        register_model_step_retry_policies: Optional[List[RetryPolicy]] = None,
         model_package_group_name=None,
         model_metrics=None,
         approval_status=None,
@@ -85,7 +85,7 @@ def __init__(
         display_name=None,
         description=None,
         tags=None,
-        model: Union[Model, PipelineModel] = None,
+        model: Optional[Union[Model, PipelineModel]] = None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
         domain=None,
@@ -328,8 +328,8 @@ def __init__(
         instance_count,
         instance_type,
         transform_inputs,
-        description: str = None,
-        display_name: str = None,
+        description: Optional[str] = None,
+        display_name: Optional[str] = None,
         # model arguments
         image_uri=None,
         predictor_cls=None,
@@ -346,9 +346,9 @@ def __init__(
         volume_kms_key=None,
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
         # step retry policies
-        repack_model_step_retry_policies: List[RetryPolicy] = None,
-        model_step_retry_policies: List[RetryPolicy] = None,
-        transform_step_retry_policies: List[RetryPolicy] = None,
+        repack_model_step_retry_policies: Optional[List[RetryPolicy]] = None,
+        model_step_retry_policies: Optional[List[RetryPolicy]] = None,
+        transform_step_retry_policies: Optional[List[RetryPolicy]] = None,
         **kwargs,
     ):
         """Construct steps required for a Transformer step collection:
diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py
@@ -362,10 +362,10 @@ def __init__(
         self,
         name: str,
         step_type: StepTypeEnum,
-        display_name: str = None,
-        description: str = None,
-        depends_on: Optional[List[Union[str, Step, "StepCollection", StepOutput]]] = None,
-        retry_policies: List[RetryPolicy] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
     ):
         super().__init__(
             name=name,
@@ -404,14 +404,14 @@ class TrainingStep(ConfigurableRetryStep):
     def __init__(
         self,
         name: str,
-        step_args: _JobStepArguments = None,
-        estimator: EstimatorBase = None,
-        display_name: str = None,
-        description: str = None,
-        inputs: Union[TrainingInput, dict, str, FileSystemInput] = None,
-        cache_config: CacheConfig = None,
+        step_args: Optional[_JobStepArguments] = None,
+        estimator: Optional[EstimatorBase] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        inputs: Optional[Union[TrainingInput, dict, str, FileSystemInput]] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
-        retry_policies: List[RetryPolicy] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
     ):
         """Construct a `TrainingStep`, given an `EstimatorBase` instance.
 
@@ -681,14 +681,14 @@ class TransformStep(ConfigurableRetryStep):
     def __init__(
         self,
         name: str,
-        step_args: _JobStepArguments = None,
-        transformer: Transformer = None,
-        inputs: TransformInput = None,
-        display_name: str = None,
-        description: str = None,
-        cache_config: CacheConfig = None,
+        step_args: Optional[_JobStepArguments] = None,
+        transformer: Optional[Transformer] = None,
+        inputs: Optional[TransformInput] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
-        retry_policies: List[RetryPolicy] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
     ):
         """Constructs a `TransformStep`, given a `Transformer` instance.
 
@@ -808,19 +808,19 @@ class ProcessingStep(ConfigurableRetryStep):
     def __init__(
         self,
         name: str,
-        step_args: _JobStepArguments = None,
-        processor: Processor = None,
-        display_name: str = None,
-        description: str = None,
-        inputs: List[ProcessingInput] = None,
-        outputs: List[ProcessingOutput] = None,
-        job_arguments: List[str] = None,
-        code: str = None,
-        property_files: List[PropertyFile] = None,
-        cache_config: CacheConfig = None,
+        step_args: Optional[_JobStepArguments] = None,
+        processor: Optional[Processor] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        inputs: Optional[List[ProcessingInput]] = None,
+        outputs: Optional[List[ProcessingOutput]] = None,
+        job_arguments: Optional[List[str]] = None,
+        code: Optional[str] = None,
+        property_files: Optional[List[PropertyFile]] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
-        retry_policies: List[RetryPolicy] = None,
-        kms_key=None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
+        kms_key: Optional[str] = None,
     ):
         """Construct a `ProcessingStep`, given a `Processor` instance.
 
@@ -980,15 +980,15 @@ class TuningStep(ConfigurableRetryStep):
     def __init__(
         self,
         name: str,
-        step_args: _JobStepArguments = None,
-        tuner: HyperparameterTuner = None,
-        display_name: str = None,
-        description: str = None,
+        step_args: Optional[_JobStepArguments] = None,
+        tuner: Optional[HyperparameterTuner] = None,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
         inputs=None,
-        job_arguments: List[str] = None,
-        cache_config: CacheConfig = None,
+        job_arguments: Optional[List[str]] = None,
+        cache_config: Optional[CacheConfig] = None,
         depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
-        retry_policies: List[RetryPolicy] = None,
+        retry_policies: Optional[List[RetryPolicy]] = None,
     ):
         """Construct a `TuningStep`, given a `HyperparameterTuner` instance.
 
