feat: exposes v1 enhanced types and adds tests

google/cloud/aiplatform/gapic/schema/__init__.py

@@ -15,31 +15,49 @@
 # limitations under the License.
 
 from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema import predict
-from google.cloud.aiplatform.v1beta1.schema import trainingjob
+from google.cloud.aiplatform.v1.schema import predict
+from google.cloud.aiplatform.v1.schema import trainingjob
+from google.cloud.aiplatform.v1beta1.schema import predict as predict_v1beta1
+from google.cloud.aiplatform.v1beta1.schema import predict as trainingjob_v1beta1
+
+# import the v1 submodules for enhancement
+from google.cloud.aiplatform.v1.schema.predict.instance_v1 import types as instance
+from google.cloud.aiplatform.v1.schema.predict.params_v1 import types as params
+from google.cloud.aiplatform.v1.schema.predict.prediction_v1 import types as prediction
+from google.cloud.aiplatform.v1.schema.trainingjob.definition_v1 import (
+    types as definition,
+)
+
+# import the v1beta1 submodules for enhancement
 from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1 import (
-    types as instance,
+    types as instance_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import (
-    types as params,
+    types as params_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
-    types as prediction,
+    types as prediction_v1beta1,
 )
 from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
-    types as definition,
+    types as definition_v1beta1,
 )
 
 __all__ = (
     "predict",
     "trainingjob",
+    "predict_v1beta1",
+    "trainingjob_v1beta1",
 )
 
 enhanced_types_packages = [
     instance,
     params,
     prediction,
     definition,
+    instance_v1beta1,
+    params_v1beta1,
+    prediction_v1beta1,
+    definition_v1beta1,
 ]
 
 for pkg in enhanced_types_packages:

google/cloud/aiplatform/v1/schema/__init__.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema import predict
+from google.cloud.aiplatform.v1.schema import trainingjob
+
+
+__all__ = (
+    "predict",
+    "trainingjob",
+)

google/cloud/aiplatform/v1/schema/predict/__init__.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema.predict import instance
+from google.cloud.aiplatform.v1.schema.predict import params
+from google.cloud.aiplatform.v1.schema.predict import prediction
+
+__all__ = (
+    "instance",
+    "params",
+    "prediction",
+)

google/cloud/aiplatform/v1/schema/trainingjob/__init__.py

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.cloud.aiplatform.v1.schema.trainingjob import definition
+
+__all__ = ("definition",)

tests/unit/enhanced_library/test_enhanced_types.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 from __future__ import absolute_import
 
-from google.cloud.aiplatform.v1beta1.schema.trainingjob import definition
+from google.cloud.aiplatform.v1.schema.trainingjob import definition
+from google.cloud.aiplatform.v1beta1.schema.trainingjob import (
+    definition as definition_v1beta1,
+)
 
 ModelType = definition.AutoMlImageClassificationInputs().ModelType
 test_training_input = definition.AutoMlImageClassificationInputs(
@@ -23,14 +26,36 @@
     disable_early_stopping=False,
 )
 
+ModelType_v1beta1 = definition_v1beta1.AutoMlImageClassificationInputs().ModelType
+test_training_input_v1beta1 = definition.AutoMlImageClassificationInputs(
+    multi_label=True,
+    model_type=ModelType_v1beta1.CLOUD,
+    budget_milli_node_hours=8000,
+    disable_early_stopping=False,
+)
+
 
-def test_exposes_to_value_method():
+# Test the v1 enhanced types.
+def test_exposes_to_value_method_v1():
     assert hasattr(test_training_input, "to_value")
 
 
-def test_exposes_from_value_method():
+def test_exposes_from_value_method_v1():
     assert hasattr(test_training_input, "from_value")
 
 
-def test_exposes_from_map_method():
+def test_exposes_from_map_method_v1():
     assert hasattr(test_training_input, "from_map")
+
+
+# Test the v1beta1 enhanced types.
+def test_exposes_to_value_method_v1beta1():
+    assert hasattr(test_training_input_v1beta1, "to_value")
+
+
+def test_exposes_from_value_method_v1beta1():
+    assert hasattr(test_training_input_v1beta1, "from_value")
+
+
+def test_exposes_from_map_method_v1beta1():
+    assert hasattr(test_training_input_v1beta1, "from_map")