Wrong type hints for optional parameters

[AdrianB-sovo]

Describe the bug Some optional parameters don't use Union[<type>, None] = None, which is problematic for those that use a type checker.

To reproduce

from typing import List
from sagemaker.workflow.steps import ProcessingStep

def get_job_args() -> List[str] | None
    # In my use-case, the arguments depend on some config, which can be empty.
    return None

process_step = ProcessingStep(
    name="some name",
    job_arguments=get_job_args(), # <- Type checker complains that `List[str] | None` is incompatible with `List[str]`.
)

Expected behavior The type checker shouldn't give me errors when I use Sagemaker in a valid way.

System information A description of your system. Please provide:

• SageMaker Python SDK version: 2.215.0
• Python version: 3.11.7

Additional context I'm using VScode with Pylance configured in strict mode.

[liujiaorr]

Hi AdrianB, can you give a sample of correct use of it, and do you have any idea how to fix it systematically? Thanks

[AdrianB-sovo]

Every places that have defaults values of None should either take Optional[...] or Union[...]. For example, this:

class ProcessingStep(ConfigurableRetryStep):
    """`ProcessingStep` for SageMaker Pipelines Workflows."""

    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,
        depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
        retry_policies: List[RetryPolicy] = None,
        kms_key=None,
    ):

Could become:

class ProcessingStep(ConfigurableRetryStep):
    """`ProcessingStep` for SageMaker Pipelines Workflows."""

    def __init__(
        self,
        name: str,
        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: Optional[List[RetryPolicy]] = None,
        kms_key=None,
    ):

[qidewenwhen]

The fix has been merged: https://github.com/aws/sagemaker-python-sdk/blob/1b4dc7cfc4ca5b3fa3c9cfefe5a57d35449000fc/src/sagemaker/workflow/steps.py#L805-L823

Closing this issue. Feel free to reopen if you have further questions. Thanks!