From a04a5613cec36811db8768da5ea7c3229da3074b Mon Sep 17 00:00:00 2001
From: Eric Schmidt <erschmid@google.com>
Date: Mon, 8 Feb 2021 11:01:59 -0800
Subject: [PATCH] feat: moves manual enhanced lib edits outside of generated
 files (#198)

---
 .../cloud/aiplatform/gapic/schema/__init__.py | 25 +++++++++++++++++--
 .../schema/predict/instance/__init__.py       |  4 ---
 .../v1beta1/schema/predict/params/__init__.py |  4 ---
 .../schema/predict/prediction/__init__.py     |  6 -----
 .../schema/trainingjob/definition/__init__.py |  6 -----
 5 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/google/cloud/aiplatform/gapic/schema/__init__.py b/google/cloud/aiplatform/gapic/schema/__init__.py
index 01fd26b30a..8b93f6ff9a 100644
--- a/google/cloud/aiplatform/gapic/schema/__init__.py
+++ b/google/cloud/aiplatform/gapic/schema/__init__.py
@@ -14,12 +14,33 @@
 # See the License for the specific language governing permissions and
 # 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.v1beta1.schema.predict.instance_v1beta1 import (
+    types as instance,
+)
+from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import (
+    types as params,
+)
+from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
+    types as prediction,
+)
+from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
+    types as definition,
+)
 
 __all__ = (
     "predict",
     "trainingjob",
 )
+
+enhanced_types_packages = [
+    instance,
+    params,
+    prediction,
+    definition,
+]
+
+for pkg in enhanced_types_packages:
+    _decorators._add_methods_to_classes_in_package(pkg)
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/__init__.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/__init__.py
index 3c4e8af160..32ad0e9ff2 100644
--- a/google/cloud/aiplatform/v1beta1/schema/predict/instance/__init__.py
+++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance/__init__.py
@@ -14,9 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1 import types as pkg
-
 from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1.types.image_classification import (
     ImageClassificationPredictionInstance,
 )
@@ -56,4 +53,3 @@
     "VideoClassificationPredictionInstance",
     "VideoObjectTrackingPredictionInstance",
 )
-_decorators._add_methods_to_classes_in_package(pkg)
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/__init__.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/__init__.py
index 45471523c9..4a410f3904 100644
--- a/google/cloud/aiplatform/v1beta1/schema/predict/params/__init__.py
+++ b/google/cloud/aiplatform/v1beta1/schema/predict/params/__init__.py
@@ -14,9 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import types as pkg
-
 from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1.types.image_classification import (
     ImageClassificationPredictionParams,
 )
@@ -44,4 +41,3 @@
     "VideoClassificationPredictionParams",
     "VideoObjectTrackingPredictionParams",
 )
-_decorators._add_methods_to_classes_in_package(pkg)
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/__init__.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/__init__.py
index f8a4d63d58..159824217b 100644
--- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/__init__.py
+++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/__init__.py
@@ -14,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
-    types as pkg,
-)
-
 from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1.types.classification import (
     ClassificationPredictionResult,
 )
@@ -66,4 +61,3 @@
     "VideoClassificationPredictionResult",
     "VideoObjectTrackingPredictionResult",
 )
-_decorators._add_methods_to_classes_in_package(pkg)
diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/__init__.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/__init__.py
index 8c71f1f7cf..392fae649e 100644
--- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/__init__.py
+++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/__init__.py
@@ -14,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
-    types as pkg,
-)
-
 from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1.types.automl_forecasting import (
     AutoMlForecasting,
 )
@@ -134,4 +129,3 @@
     "AutoMlVideoObjectTrackingInputs",
     "ExportEvaluatedDataItemsConfig",
 )
-_decorators._add_methods_to_classes_in_package(pkg)