[AML-96242] Add Type annotations

Author: stisac

src/sagemaker/lineage/_utils.py

@@ -45,4 +45,4 @@ def get_module(module_name):
     Returns:
         [obj]: The imported module
     """
-    return import_module(module_name)
+    return import_module(module_name)

src/sagemaker/lineage/action.py

@@ -13,8 +13,13 @@
 """This module contains code to create and manage SageMaker ``Actions``."""
 from __future__ import absolute_import
 
+from typing import Type, Iterator, Optional
+from datetime import datetime
+
+from sagemaker import Session
 from sagemaker.apiutils import _base_types
 from sagemaker.lineage import _api_types, _utils
+from sagemaker.lineage._api_types import ActionSource
 
 
 class Action(_base_types.Record):
@@ -53,46 +58,46 @@ class Action(_base_types.Record):
         last_modified_by (obj): Contextual info on which account created the action.
     """
 
-    action_arn = None
-    action_name = None
-    action_type = None
-    description = None
-    status = None
-    source = None
-    properties = None
-    properties_to_remove = None
-    tags = None
-    creation_time = None
-    created_by = None
-    last_modified_time = None
-    last_modified_by = None
-
-    _boto_create_method = "create_action"
-    _boto_load_method = "describe_action"
-    _boto_update_method = "update_action"
-    _boto_delete_method = "delete_action"
-
-    _boto_update_members = [
+    action_arn: str = None
+    action_name: str = None
+    action_type: str = None
+    description: str = None
+    status: str = None
+    source: object = None
+    properties: dict[str, str] = None
+    properties_to_remove: list[str] = None
+    tags: list[dict[str, str]] = None
+    creation_time: datetime = None
+    created_by: object = None
+    last_modified_time: datetime = None
+    last_modified_by: object = None
+
+    _boto_create_method: str = "create_action"
+    _boto_load_method: str = "describe_action"
+    _boto_update_method: str = "update_action"
+    _boto_delete_method: str = "delete_action"
+
+    _boto_update_members: list[str] = [
         "action_name",
         "description",
         "status",
         "properties",
         "properties_to_remove",
     ]
 
-    _boto_delete_members = ["action_name"]
+    _boto_delete_members: list[str] = ["action_name"]
 
-    _custom_boto_types = {"source": (_api_types.ActionSource, False)}
+    _custom_boto_types: dict[str, tuple[Type[ActionSource], bool]] = {"source": (_api_types.ActionSource, False)}
 
-    def save(self):
+    def save(self) -> object:
         """Save the state of this Action to SageMaker.
 
         Returns:
             Action: A SageMaker ``Action``object.
         """
         return self._invoke_api(self._boto_update_method, self._boto_update_members)
 
-    def delete(self, disassociate=False):
+    def delete(self, disassociate: bool = False):
         """Delete the action.
 
         Args:
@@ -110,7 +115,7 @@ def delete(self, disassociate=False):
         self._invoke_api(self._boto_delete_method, self._boto_delete_members)
 
     @classmethod
-    def load(cls, action_name, sagemaker_session=None):
+    def load(cls, action_name: str, sagemaker_session: Session = None) -> object:
         """Load an existing action and return an ``Action`` object representing it.
 
         Args:
@@ -130,7 +135,7 @@ def load(cls, action_name, sagemaker_session=None):
         )
         return result
 
-    def set_tag(self, tag=None):
+    def set_tag(self, tag: object = None) -> list[{str: str}]:
         """Add a tag to the object.
 
         Args:
@@ -140,7 +145,7 @@ def set_tag(self, tag=None):
         """
         return self._set_tags(resource_arn=self.action_arn, tags=[tag])
 
-    def set_tags(self, tags=None):
+    def set_tags(self, tags: list[{str: str}] = None) -> list[{str: str}]:
         """Add tags to the object.
 
         Args:
@@ -153,17 +158,17 @@ def set_tags(self, tags=None):
 
     @classmethod
     def create(
-        cls,
-        action_name=None,
-        source_uri=None,
-        source_type=None,
-        action_type=None,
-        description=None,
-        status=None,
-        properties=None,
-        tags=None,
-        sagemaker_session=None,
-    ):
+            cls,
+            action_name: str = None,
+            source_uri: str = None,
+            source_type: str = None,
+            action_type: str = None,
+            description: str = None,
+            status: str = None,
+            properties: dict[str:str] = None,
+            tags: list[dict[str:str]] = None,
+            sagemaker_session: Session = None,
+    ) -> object:
         """Create an action and return an ``Action`` object representing it.
 
         Args:
@@ -197,17 +202,17 @@ def create(
 
     @classmethod
     def list(
-        cls,
-        source_uri=None,
-        action_type=None,
-        created_after=None,
-        created_before=None,
-        sort_by=None,
-        sort_order=None,
-        sagemaker_session=None,
-        max_results=None,
-        next_token=None,
-    ):
+            cls,
+            source_uri: Optional[str] = None,
+            action_type: Optional[str] = None,
+            created_after: Optional[datetime] = None,
+            created_before: Optional[datetime] = None,
+            sort_by: Optional[str] = None,
+            sort_order: Optional[str] = None,
+            sagemaker_session: Session = None,
+            max_results: Optional[int] = None,
+            next_token: Optional[str] = None,
+    ) -> Iterator:
         """Return a list of action summaries.
 
         Args: