Given options for both uri and (namespace, bucket) in add_artifact, remove_artifact methods

Author: YashPandit4u

ads/model/datascience_model.py

@@ -1470,8 +1470,9 @@ def _download_file_description_artifact(self) -> Tuple[Union[str, List[str]], in
 
     def add_artifact(
         self,
-        namespace: str,
-        bucket: str,
+        uri: Optional[str] = None,
+        namespace: Optional[str] = None,
+        bucket: Optional[str] = None,
         prefix: Optional[str] = None,
         files: Optional[List[str]] = None,
     ):
@@ -1499,6 +1500,16 @@ def add_artifact(
         - If no objects are found to add to the model description, a ValueError is raised.
         """
 
+        if uri and (namespace or bucket):
+            raise ValueError("Either 'uri' must be provided or both 'namespace' and 'bucket' must be provided.")
+        if uri:
+            object_storage_details = ObjectStorageDetails.from_path(uri)
+            bucket = object_storage_details.bucket
+            namespace = object_storage_details.namespace
+            prefix = None if object_storage_details.filepath == "" else object_storage_details.filepath
+        if ((not namespace) or (not bucket)):
+            raise ValueError("Both 'namespace' and 'bucket' must be provided.")
+
         # Check if both prefix and files are provided
         if prefix is not None and files is not None:
             raise ValueError("Both 'prefix' and 'files' cannot be provided. Please provide only one.")
@@ -1590,7 +1601,13 @@ def list_obj_versions_unpaginated():
         )
         self.set_spec(self.CONST_MODEL_FILE_DESCRIPTION, tmp_model_file_description)
 
-    def remove_artifact(self, namespace: str, bucket: str, prefix: Optional[str] = None):
+    def remove_artifact(
+            self, 
+            uri: Optional[str] = None, 
+            namespace: Optional[str] = None, 
+            bucket: Optional[str] = None, 
+            prefix: Optional[str] = None
+        ):
         """
         Removes information about objects in a specified bucket from the model description JSON.
 
@@ -1607,6 +1624,16 @@ def remove_artifact(self, namespace: str, bucket: str, prefix: Optional[str] = N
         ValueError: If the model description JSON is None.
         """
 
+        if uri and (namespace or bucket):
+            raise ValueError("Either 'uri' must be provided or both 'namespace' and 'bucket' must be provided.")
+        if uri:
+            object_storage_details = ObjectStorageDetails.from_path(uri)
+            bucket = object_storage_details.bucket
+            namespace = object_storage_details.namespace
+            prefix = None if object_storage_details.filepath == "" else object_storage_details.filepath
+        if ((not namespace) or (not bucket)):
+            raise ValueError("Both 'namespace' and 'bucket' must be provided.")
+
         def findModelIdx():
             for idx, model in enumerate(self.model_file_description["models"]):
                 if (