From b478075efb05553760514256fee9a63126a9916f Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Mon, 9 Aug 2021 08:52:56 -0700 Subject: [PATCH] feat: Added the VertexAiResourceNoun.to_dict() method (#588) --- google/cloud/aiplatform/base.py | 5 +++++ google/cloud/aiplatform/metadata/resource.py | 3 +-- tests/system/aiplatform/test_dataset.py | 6 ++---- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 1a3eed8add..20f9aa07ad 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -44,6 +44,7 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec +from google.protobuf import json_format logging.basicConfig(level=logging.INFO, stream=sys.stdout) @@ -607,6 +608,10 @@ def _assert_gca_resource_is_available(self) -> None: def __repr__(self) -> str: return f"{object.__repr__(self)} \nresource name: {self.resource_name}" + def to_dict(self) -> Dict[str, Any]: + """Returns the resource proto as a dictionary.""" + return json_format.MessageToDict(self.gca_resource._pb) + def optional_sync( construct_object_on_arg: Optional[str] = None, diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 85ac419d40..37bb7327cd 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -24,7 +24,6 @@ import proto from google.api_core import exceptions from google.auth import credentials as auth_credentials -from google.protobuf import json_format from google.cloud.aiplatform import base, initializer from google.cloud.aiplatform import utils @@ -98,7 +97,7 @@ def __init__( @property def metadata(self) -> Dict: - return json_format.MessageToDict(self._gca_resource._pb)["metadata"] + return self.to_dict()["metadata"] @property def schema_title(self) -> str: diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py index e852933dc3..0167cb8f20 100644 --- a/tests/system/aiplatform/test_dataset.py +++ b/tests/system/aiplatform/test_dataset.py @@ -21,7 +21,6 @@ import importlib from google import auth as google_auth -from google.protobuf import json_format from google.api_core import exceptions from google.api_core import client_options @@ -239,16 +238,15 @@ def test_create_tabular_dataset(self, dataset_gapic_client, shared_state): gcs_source=[_TEST_TABULAR_CLASSIFICATION_GCS_SOURCE], ) - gapic_dataset = tabular_dataset._gca_resource shared_state["dataset_name"] = tabular_dataset.resource_name - gapic_metadata = json_format.MessageToDict(gapic_dataset._pb.metadata) + gapic_metadata = tabular_dataset.to_dict()["metadata"] gcs_source_uris = gapic_metadata["inputConfig"]["gcsSource"]["uri"] assert len(gcs_source_uris) == 1 assert _TEST_TABULAR_CLASSIFICATION_GCS_SOURCE == gcs_source_uris[0] assert ( - gapic_dataset.metadata_schema_uri + tabular_dataset.metadata_schema_uri == aiplatform.schema.dataset.metadata.tabular )