feat: Make aiplatform.Dataset private (#296)

- Removes all use of `aiplatform.Dataset` - Replaces internal references with `aiplatform._Dataset` - Replaces public references with specific dataset sub-class or TypeVar `datasets.SupportedDatasets` ## TODO after merging - [x] @sasha-gitg PTAL [at this list of updates](https://docs.google.com/document/d/1mKqldyCZaLgtYrgxpd6F3TJ35Ap3l6MhOxyiWTmZ4NQ/edit) to external notebooks and Colabs to no longer reference `aiplatform.Dataset` Fixes [b/184156060](http://b/184156060) 🦕
googleapis · Apr 11, 2021 · 1f0d5f3 · 1f0d5f3
1 parent 3ec9386
commit 1f0d5f3
Show file tree

Hide file tree

Showing 16 changed files with 168 additions and 90 deletions.
diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py
@@ -20,9 +20,8 @@
 
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform.datasets import (
-    Dataset,
-    TabularDataset,
     ImageDataset,
+    TabularDataset,
     TextDataset,
     VideoDataset,
 )
@@ -59,7 +58,6 @@
     "CustomTrainingJob",
     "CustomContainerTrainingJob",
     "CustomPythonPackageTrainingJob",
-    "Dataset",
     "Endpoint",
     "ImageDataset",
     "Model",

diff --git a/google/cloud/aiplatform/datasets/__init__.py b/google/cloud/aiplatform/datasets/__init__.py
@@ -15,14 +15,15 @@
 # limitations under the License.
 #
 
-from google.cloud.aiplatform.datasets.dataset import Dataset
+from google.cloud.aiplatform.datasets.dataset import _Dataset
 from google.cloud.aiplatform.datasets.tabular_dataset import TabularDataset
 from google.cloud.aiplatform.datasets.image_dataset import ImageDataset
 from google.cloud.aiplatform.datasets.text_dataset import TextDataset
 from google.cloud.aiplatform.datasets.video_dataset import VideoDataset
 
+
 __all__ = (
-    "Dataset",
+    "_Dataset",
     "TabularDataset",
     "ImageDataset",
     "TextDataset",

diff --git a/google/cloud/aiplatform/datasets/dataset.py b/google/cloud/aiplatform/datasets/dataset.py
@@ -33,7 +33,7 @@
 from google.cloud.aiplatform.datasets import _datasources
 
 
-class Dataset(base.AiPlatformResourceNounWithFutureManager):
+class _Dataset(base.AiPlatformResourceNounWithFutureManager):
     """Managed dataset resource for AI Platform"""
 
     client_class = utils.DatasetClientWithOverride
@@ -115,7 +115,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
-    ) -> "Dataset":
+    ) -> "_Dataset":
         """Creates a new dataset and optionally imports data into dataset when
         source and import_schema_uri are passed.
 
@@ -241,7 +241,7 @@ def _create_and_import(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
         sync: bool = True,
-    ) -> "Dataset":
+    ) -> "_Dataset":
         """Creates a new dataset and optionally imports data into dataset when
         source and import_schema_uri are passed.
 
@@ -400,7 +400,7 @@ def import_data(
         import_schema_uri: str,
         data_item_labels: Optional[Dict] = None,
         sync: bool = True,
-    ) -> "Dataset":
+    ) -> "_Dataset":
         """Upload data to existing managed dataset.
 
         Args:

diff --git a/google/cloud/aiplatform/datasets/image_dataset.py b/google/cloud/aiplatform/datasets/image_dataset.py
@@ -26,7 +26,7 @@
 from google.cloud.aiplatform import utils
 
 
-class ImageDataset(datasets.Dataset):
+class ImageDataset(datasets._Dataset):
     """Managed image dataset resource for AI Platform"""
 
     _supported_metadata_schema_uris: Optional[Tuple[str]] = (

diff --git a/google/cloud/aiplatform/datasets/tabular_dataset.py b/google/cloud/aiplatform/datasets/tabular_dataset.py
@@ -26,7 +26,7 @@
 from google.cloud.aiplatform import utils
 
 
-class TabularDataset(datasets.Dataset):
+class TabularDataset(datasets._Dataset):
     """Managed tabular dataset resource for AI Platform"""
 
     _supported_metadata_schema_uris: Optional[Tuple[str]] = (

diff --git a/google/cloud/aiplatform/datasets/text_dataset.py b/google/cloud/aiplatform/datasets/text_dataset.py
@@ -26,7 +26,7 @@
 from google.cloud.aiplatform import utils
 
 
-class TextDataset(datasets.Dataset):
+class TextDataset(datasets._Dataset):
     """Managed text dataset resource for AI Platform"""
 
     _supported_metadata_schema_uris: Optional[Tuple[str]] = (

diff --git a/google/cloud/aiplatform/datasets/video_dataset.py b/google/cloud/aiplatform/datasets/video_dataset.py
@@ -26,7 +26,7 @@
 from google.cloud.aiplatform import utils
 
 
-class VideoDataset(datasets.Dataset):
+class VideoDataset(datasets._Dataset):
     """Managed video dataset resource for AI Platform"""
 
     _supported_metadata_schema_uris: Optional[Tuple[str]] = (