next-mypy fixes azure-ai-ml sdk (#44132)

* next-mypy fixes for automl in azure-ai-ml sdk

* remainging next mypy issues fix

* minimizing the changes

* resolving type issue with control flow node file

* minimizing the mypy changes

* minimizing the mypy changes

* minimizing the mypy changes

* minimizing the mypy changes

---------

Co-authored-by: Zubairuddin Mohammed <[email protected]>

Author: ZubaeyrMSFT

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],
     ) -> None:
         """Sets the interactive services for the node.
 

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

@@ -166,5 +166,7 @@ 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)
-        load_from_rest_obj_func = pipeline_node_factory.get_load_from_rest_object_func(_type=node_type)
+        load_from_rest_obj_func = pipeline_node_factory.get_load_from_rest_object_func(
+            _type=node_type  # type: ignore[arg-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: