next-mypy fixes azure-ai-ml sdk

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py

@@ -140,11 +140,11 @@ def jars(self) -> Optional[List[str]]:
         return self._jars
 
     @jars.setter
-    def jars(self, value: Union[str, List[str]]):
+    def jars(self, value: Optional[Union[str, List[str]]]):
         """Set the jars of the component.
 
         :param value: The jars of the component.
-        :type value: Union[str, List[str]]
+        :type value: Optional[Union[str, List[str]]]
         :return: No return
         :rtype: None
         """

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/data.py

@@ -120,7 +120,7 @@ def path(self) -> Optional[Union[Path, str, PathLike]]:
         return self._path
 
     @path.setter
-    def path(self, value: str) -> None:
+    def path(self, value: Optional[Union[str, PathLike]]) -> None:
         # Call the parent setter to resolve the path with base_path if it was a local path
         # TODO: Bug Item number: 2883424
         super(Data, type(self)).path.fset(self, value)  # type: ignore

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/asset.py

@@ -85,11 +85,11 @@ def version(self) -> Optional[str]:
         return self._version
 
     @version.setter
-    def version(self, value: str) -> None:
+    def version(self, value: Optional[str]) -> None:
         """Sets the asset version.
 
         :param value: The asset version.
-        :type value: str
+        :type value: Optional[str]
         :raises ValidationException: Raised if value is not a string.
         """
         if value:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/environment.py

@@ -284,7 +284,7 @@ def _from_container_rest_object(cls, env_container_rest_object: EnvironmentConta
 
         # Setting version to None since if version is not provided it is defaulted to "1".
         # This should go away once container concept is finalized.
-        env.version = None
+        env.version = None  # type: ignore[assignment]
         return env
 
     def _to_arm_resource_param(self, **kwargs: Any) -> Dict:  # pylint: disable=unused-argument
@@ -404,7 +404,7 @@ def _localize(self, base_path: str) -> None:
         self._creation_context = None
         self._base_path = base_path
         if self._is_anonymous:
-            self.name, self.version = None, None
+            self.name, self.version = None, None  # type: ignore[assignment]
 
 
 # TODO: Remove _DockerBuild and _DockerConfiguration classes once local endpoint moves to using updated env

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/base_node.py

@@ -205,11 +205,11 @@ def name(self) -> Optional[str]:
         return self._name
 
     @name.setter
-    def name(self, value: str) -> None:
+    def name(self, value: Optional[str]) -> None:
         """Set the name of the node.
 
         :param value: The name to set for the node.
-        :type value: str
+        :type value: Optional[str]
         :return: None
         """
         # when name is not lower case, lower it to make sure it's a valid node name

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py

@@ -260,13 +260,22 @@ def distribution(
     @distribution.setter
     def distribution(
         self,
-        value: Union[Dict, PyTorchDistribution, TensorFlowDistribution, MpiDistribution, RayDistribution],
+        value: Optional[
+            Union[
+                Dict,
+                PyTorchDistribution,
+                TensorFlowDistribution,
+                MpiDistribution,
+                RayDistribution,
+                DistributionConfiguration,
+            ]
+        ],
     ) -> None:
         """Sets the configuration for the distributed command component or job.
 
         :param value: The configuration for distributed jobs.
-        :type value: Union[dict, ~azure.ai.ml.PyTorchDistribution, ~azure.ai.ml.MpiDistribution,
-            ~azure.ai.ml.TensorFlowDistribution, ~azure.ai.ml.RayDistribution]
+        :type value: Optional[Union[dict, ~azure.ai.ml.PyTorchDistribution, ~azure.ai.ml.MpiDistribution,
+            ~azure.ai.ml.TensorFlowDistribution, ~azure.ai.ml.RayDistribution]]
         """
         if isinstance(value, dict):
             dist_schema = UnionField(
@@ -309,11 +318,11 @@ def queue_settings(self) -> Optional[QueueSettings]:
         return self._queue_settings
 
     @queue_settings.setter
-    def queue_settings(self, value: Union[Dict, QueueSettings]) -> None:
+    def queue_settings(self, value: Optional[Union[Dict, QueueSettings]]) -> None:
         """Sets the queue settings for the command component or job.
 
         :param value: The queue settings for the command component or job.
-        :type value: Union[dict, ~azure.ai.ml.entities.QueueSettings]
+        :type value: Optional[Union[dict, ~azure.ai.ml.entities.QueueSettings]]
         """
         if isinstance(value, dict):
             value = QueueSettings(**value)
@@ -373,7 +382,7 @@ def services(
     @services.setter
     def services(
         self,
-        value: Dict,
+        value: Optional[Dict[str, Any]],
     ) -> None:
         """Sets the interactive services for the node.
 
@@ -951,7 +960,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> "Command":
             node.distribution = copy.deepcopy(self.distribution)
             node.resources = copy.deepcopy(self.resources)
             node.queue_settings = copy.deepcopy(self.queue_settings)
-            node.services = copy.deepcopy(self.services)
+            node.services = copy.deepcopy(self.services) if self.services is not None else None
             node.identity = copy.deepcopy(self.identity)
             return node
         msg = "Command can be called as a function only when referenced component is {}, currently got {}."

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/control_flow_node.py

@@ -165,6 +165,6 @@ def _is_loop_node_dict(obj: Any) -> bool:
     def _from_rest_object(cls, obj: dict, pipeline_jobs: dict) -> "LoopNode":
         from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
 
-        node_type = obj.get(CommonYamlFields.TYPE, None)
+        node_type = cast(str, obj.get(CommonYamlFields.TYPE, None))
         load_from_rest_obj_func = pipeline_node_factory.get_load_from_rest_object_func(_type=node_type)
         return load_from_rest_obj_func(obj, pipeline_jobs)  # type: ignore

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/data_transfer.py

@@ -493,7 +493,7 @@ def sink(self) -> Optional[Union[Dict, Database, FileSystem]]:
         return self._sink
 
     @sink.setter
-    def sink(self, value: Union[Dict, Database, FileSystem]) -> None:
+    def sink(self, value: Optional[Union[Dict, Database, FileSystem]]) -> None:
         self._sink = _build_source_sink(value)
 
     @classmethod

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py

@@ -252,11 +252,12 @@ def resources(self) -> Optional[JobResourceConfiguration]:
         return self._resources
 
     @resources.setter
-    def resources(self, value: Union[JobResourceConfiguration, Dict]) -> None:
+    def resources(self, value: Optional[Union[JobResourceConfiguration, Dict]]) -> None:
         """Set the resource configuration for the parallel job.
 
         :param value: The resource configuration for the parallel job.
-        :type value: ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration or dict
+        :type value: Optional[Union[
+            ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration, dict]]
         """
         if isinstance(value, dict):
             value = JobResourceConfiguration(**value)
@@ -316,11 +317,11 @@ def task(self) -> Optional[ParallelTask]:
         return self._task  # type: ignore
 
     @task.setter
-    def task(self, value: Union[ParallelTask, Dict]) -> None:
+    def task(self, value: Optional[Union[ParallelTask, Dict]]) -> None:
         """Set the parallel task.
 
         :param value: The parallel task.
-        :type value: ~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask or dict
+        :type value: Optional[Union[~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask, dict]]
         """
         # base path should be reset if task is set via sdk
         self._base_path: Optional[Union[str, os.PathLike]] = None
@@ -360,16 +361,18 @@ def set_resources(
         if self.resources is None:
             self.resources = JobResourceConfiguration()
 
+        # Use local variable for type narrowing
+        resources = self.resources
         if instance_type is not None:
-            self.resources.instance_type = instance_type
+            resources.instance_type = instance_type
         if instance_count is not None:
-            self.resources.instance_count = instance_count
+            resources.instance_count = instance_count
         if properties is not None:
-            self.resources.properties = properties
+            resources.properties = properties
         if docker_args is not None:
-            self.resources.docker_args = docker_args
+            resources.docker_args = docker_args
         if shm_size is not None:
-            self.resources.shm_size = shm_size
+            resources.shm_size = shm_size
 
         # Save the resources to internal component as well, otherwise calling sweep() will loose the settings
         if isinstance(self.component, Component):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/pipeline.py

@@ -103,11 +103,11 @@ def settings(self) -> Optional[PipelineJobSettings]:
         return self._settings
 
     @settings.setter
-    def settings(self, value: Union[PipelineJobSettings, Dict]) -> None:
+    def settings(self, value: Optional[Union[PipelineJobSettings, Dict]]) -> None:
         """Set the settings of the pipeline.
 
         :param value: The settings of the pipeline.
-        :type value: Union[~azure.ai.ml.entities.PipelineJobSettings, dict]
+        :type value: Optional[Union[~azure.ai.ml.entities.PipelineJobSettings, dict]]
         :raises TypeError: If the value is not an instance of PipelineJobSettings or a dict.
         """
         if value is not None:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/spark.py

@@ -292,13 +292,15 @@ def identity(
     @identity.setter
     def identity(
         self,
-        value: Union[Dict[str, str], ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration],
+        value: Optional[
+            Union[Dict[str, str], ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration]
+        ],
     ) -> None:
         """Sets the identity that the Spark job will use while running on compute.
 
         :param value: The identity that the Spark job will use while running on compute.
-        :type value: Union[Dict[str, str], ~azure.ai.ml.entities.ManagedIdentityConfiguration,
-            ~azure.ai.ml.entities.AmlTokenConfiguration, ~azure.ai.ml.entities.UserIdentityConfiguration]
+        :type value: Optional[Union[Dict[str, str], ~azure.ai.ml.entities.ManagedIdentityConfiguration,
+            ~azure.ai.ml.entities.AmlTokenConfiguration, ~azure.ai.ml.entities.UserIdentityConfiguration]]
         """
         if isinstance(value, dict):
             identify_schema = UnionField(

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py

@@ -158,7 +158,8 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         # TODO: get rid of self._job_inputs, self._job_outputs once we have general Input
-        self._job_inputs, self._job_outputs = inputs, outputs
+        self._job_inputs: Dict[str, Any] = inputs or {}
+        self._job_outputs: Dict[Any, Any] = outputs or {}
 
         kwargs.pop("type", None)
         BaseNode.__init__(
@@ -216,17 +217,20 @@ def search_space(
         return self._search_space
 
     @search_space.setter
-    def search_space(self, values: Dict[str, Dict[str, Union[str, int, float, dict]]]) -> None:
+    def search_space(self, values: Optional[Dict[str, Dict[str, Union[str, int, float, dict]]]]) -> None:
         """Sets the search space for the sweep job.
 
         :param values: The search space to set.
-        :type values: Dict[str, Dict[str, Union[str, int, float, dict]]]
+        :type values: Optional[Dict[str, Dict[str, Union[str, int, float, dict]]]]
         """
-        search_space: Dict = {}
-        for name, value in values.items():
-            # If value is a SearchSpace object, directly pass it to job.search_space[name]
-            search_space[name] = self._value_type_to_class(value) if isinstance(value, dict) else value
-        self._search_space = search_space
+        if values is None:
+            self._search_space = None
+        else:
+            search_space: Dict = {}
+            for name, value in values.items():
+                # If value is a SearchSpace object, directly pass it to job.search_space[name]
+                search_space[name] = self._value_type_to_class(value) if isinstance(value, dict) else value
+            self._search_space = search_space
 
     @classmethod
     def _value_type_to_class(cls, value: Any) -> Dict:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/command_component.py

@@ -188,11 +188,11 @@ def instance_count(self) -> Optional[int]:
         return self.resources.instance_count if self.resources and not isinstance(self.resources, dict) else None
 
     @instance_count.setter
-    def instance_count(self, value: int) -> None:
+    def instance_count(self, value: Optional[int]) -> None:
         """Sets the number of instances or nodes to be used by the compute target.
 
         :param value: The number of instances of nodes to be used by the compute target. Defaults to 1.
-        :type value: int
+        :type value: Optional[int]
         """
         if not value:
             return

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py

@@ -230,11 +230,11 @@ def version(self) -> Optional[str]:
         return self._version
 
     @version.setter
-    def version(self, value: str) -> None:
+    def version(self, value: Optional[str]) -> None:
         """Set the version of the component.
 
         :param value: The version of the component.
-        :type value: str
+        :type value: Optional[str]
         """
         if value:
             if not isinstance(value, str):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/parallel_component.py

@@ -179,11 +179,11 @@ def instance_count(self) -> Optional[int]:
         return self.resources.instance_count if self.resources and not isinstance(self.resources, dict) else None
 
     @instance_count.setter
-    def instance_count(self, value: int) -> None:
+    def instance_count(self, value: Optional[int]) -> None:
         """Set the value of the promoted property resources.instance_count.
 
         :param value: The value to set for resources.instance_count.
-        :type value: int
+        :type value: Optional[int]
         """
         if not value:
             return
@@ -204,11 +204,11 @@ def code(self) -> Optional[str]:
         return self.task.code if self.task else None
 
     @code.setter
-    def code(self, value: str) -> None:
+    def code(self, value: Optional[str]) -> None:
         """Set the value of the promoted property task.code.
 
         :param value: The value to set for task.code.
-        :type value: str
+        :type value: Optional[str]
         """
         if not value:
             return

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/data_index_func.py

@@ -9,16 +9,15 @@
 from typing import Any, Dict, Optional, Union
 
 from azure.ai.ml._utils._experimental import experimental
-from azure.ai.ml.constants._common import AssetTypes, LegacyAssetTypes
+from azure.ai.ml.constants._common import AssetTypes, DataIndexTypes, LegacyAssetTypes
+from azure.ai.ml.constants._component import LLMRAGComponentUri
 from azure.ai.ml.entities import PipelineJob
 from azure.ai.ml.entities._builders.base_node import pipeline_node_decorator
 from azure.ai.ml.entities._credentials import ManagedIdentityConfiguration, UserIdentityConfiguration
+from azure.ai.ml.entities._indexes.entities.data_index import DataIndex
 from azure.ai.ml.entities._inputs_outputs import Input, Output
 from azure.ai.ml.entities._job.pipeline._io import PipelineInput
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
-from azure.ai.ml.constants._common import DataIndexTypes
-from azure.ai.ml.constants._component import LLMRAGComponentUri
-from azure.ai.ml.entities._indexes.entities.data_index import DataIndex
 
 SUPPORTED_INPUTS = [
     LegacyAssetTypes.PATH,
@@ -745,4 +744,4 @@ def _resolve_connection_id(ml_client, connection: Optional[str] = None) -> Optio
         connection = ml_client.connections.get(connection_name)
         if connection is None:
             return None
-    return connection.id  # type: ignore [attr-defined]
+    return connection.id  # type: ignore[attr-defined, union-attr]

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/external_data.py

@@ -72,7 +72,7 @@ def __init__(
     ) -> None:
         # As an annotation, it is not allowed to initialize the name.
         # The name will be updated by the annotated variable name.
-        self.name = None
+        self.name: Optional[str] = None
         self.type = ExternalDataType.DATABASE
         self.connection = connection
         self.query = query

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_classification_base.py

@@ -60,11 +60,11 @@ def training_parameters(self) -> Optional[ImageModelSettingsClassification]:
         return self._training_parameters
 
     @training_parameters.setter
-    def training_parameters(self, value: Union[Dict, ImageModelSettingsClassification]) -> None:
+    def training_parameters(self, value: Optional[Union[Dict, ImageModelSettingsClassification]]) -> None:
         """Setting Image training parameters for AutoML Image Classification and Image Classification Multilabel tasks.
 
         :param value: Training parameters for AutoML Image Classification and Image Classification Multilabel tasks.
-        :type value: Union[Dict, ~azure.ai.ml.automl.ImageModelSettingsClassification]
+        :type value: Optional[Union[Dict, ~azure.ai.ml.automl.ImageModelSettingsClassification]]
         :raises ~azure.ml.exceptions.ValidationException if value is not a dictionary or
          ImageModelSettingsClassification.
         :return: None

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_object_detection_base.py

@@ -64,7 +64,7 @@ def training_parameters(self) -> Optional[ImageModelSettingsObjectDetection]:
         return self._training_parameters
 
     @training_parameters.setter
-    def training_parameters(self, value: Union[Dict, ImageModelSettingsObjectDetection]) -> None:
+    def training_parameters(self, value: Optional[Union[Dict, ImageModelSettingsObjectDetection]]) -> None:
         if value is None:
             self._training_parameters = None
         elif isinstance(value, ImageModelSettingsObjectDetection):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/automl_nlp_job.py

@@ -100,7 +100,7 @@ def training_parameters(self) -> Optional[NlpFixedParameters]:
         return self._training_parameters
 
     @training_parameters.setter
-    def training_parameters(self, value: Union[Dict, NlpFixedParameters]) -> None:
+    def training_parameters(self, value: Optional[Union[Dict, NlpFixedParameters]]) -> None:
         if value is None:
             self._training_parameters = None
         elif isinstance(value, NlpFixedParameters):

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py

@@ -124,7 +124,7 @@ def training(self) -> ClassificationTrainingSettings:
         return self._training or ClassificationTrainingSettings()
 
     @training.setter
-    def training(self, value: Union[Dict, ClassificationTrainingSettings]) -> None:  # pylint: disable=unused-argument
+    def training(self, value: Union[Dict, TrainingSettings]) -> None:  # pylint: disable=unused-argument
         ...
 
     def _to_rest_object(self) -> JobBase: