feat: add context manager support in client

google/cloud/automl_v1/services/auto_ml/async_client.py

@@ -1904,6 +1904,12 @@ async def list_model_evaluations(
         # Done; return the response.
         return response
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/automl_v1/services/auto_ml/client.py

@@ -435,10 +435,7 @@ def __init__(
                 client_cert_source_for_mtls=client_cert_source_func,
                 quota_project_id=client_options.quota_project_id,
                 client_info=client_info,
-                always_use_jwt_access=(
-                    Transport == type(self).get_transport_class("grpc")
-                    or Transport == type(self).get_transport_class("grpc_asyncio")
-                ),
+                always_use_jwt_access=True,
             )
 
     def create_dataset(
@@ -2063,6 +2060,19 @@ def list_model_evaluations(
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Releases underlying transport's resources.
+
+        .. warning::
+            ONLY use as a context manager if the transport is NOT shared
+            with other clients! Exiting the with block will CLOSE the transport
+            and may cause errors in other clients!
+        """
+        self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/automl_v1/services/auto_ml/transports/base.py

@@ -324,6 +324,15 @@ def _prep_wrapped_messages(self, client_info):
             ),
         }
 
+    def close(self):
+        """Closes resources associated with the transport.
+
+       .. warning::
+            Only call this method if the transport is NOT shared
+            with other clients - this may cause errors in other clients!
+        """
+        raise NotImplementedError()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Return the client designed to process long-running operations."""

google/cloud/automl_v1/services/auto_ml/transports/grpc.py

@@ -780,5 +780,8 @@ def list_model_evaluations(
             )
         return self._stubs["list_model_evaluations"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("AutoMlGrpcTransport",)

google/cloud/automl_v1/services/auto_ml/transports/grpc_asyncio.py

@@ -794,5 +794,8 @@ def list_model_evaluations(
             )
         return self._stubs["list_model_evaluations"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("AutoMlGrpcAsyncIOTransport",)

google/cloud/automl_v1/services/prediction_service/async_client.py

@@ -544,6 +544,12 @@ async def batch_predict(
         # Done; return the response.
         return response
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/automl_v1/services/prediction_service/client.py

@@ -355,10 +355,7 @@ def __init__(
                 client_cert_source_for_mtls=client_cert_source_func,
                 quota_project_id=client_options.quota_project_id,
                 client_info=client_info,
-                always_use_jwt_access=(
-                    Transport == type(self).get_transport_class("grpc")
-                    or Transport == type(self).get_transport_class("grpc_asyncio")
-                ),
+                always_use_jwt_access=True,
             )
 
     def predict(
@@ -733,6 +730,19 @@ def batch_predict(
         # Done; return the response.
         return response
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Releases underlying transport's resources.
+
+        .. warning::
+            ONLY use as a context manager if the transport is NOT shared
+            with other clients! Exiting the with block will CLOSE the transport
+            and may cause errors in other clients!
+        """
+        self.transport.close()
+
 
 try:
     DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/automl_v1/services/prediction_service/transports/base.py

@@ -162,6 +162,15 @@ def _prep_wrapped_messages(self, client_info):
             ),
         }
 
+    def close(self):
+        """Closes resources associated with the transport.
+
+       .. warning::
+            Only call this method if the transport is NOT shared
+            with other clients - this may cause errors in other clients!
+        """
+        raise NotImplementedError()
+
     @property
     def operations_client(self) -> operations_v1.OperationsClient:
         """Return the client designed to process long-running operations."""

google/cloud/automl_v1/services/prediction_service/transports/grpc.py

@@ -354,5 +354,8 @@ def batch_predict(
             )
         return self._stubs["batch_predict"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("PredictionServiceGrpcTransport",)

google/cloud/automl_v1/services/prediction_service/transports/grpc_asyncio.py

@@ -362,5 +362,8 @@ def batch_predict(
             )
         return self._stubs["batch_predict"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("PredictionServiceGrpcAsyncIOTransport",)

google/cloud/automl_v1/types/annotation_payload.py

@@ -29,6 +29,7 @@
 
 class AnnotationPayload(proto.Message):
     r"""Contains annotation information that is relevant to AutoML.
+
     Attributes:
         translation (google.cloud.automl_v1.types.TranslationAnnotation):
             Annotation details for translation.

google/cloud/automl_v1/types/annotation_spec.py

@@ -23,6 +23,7 @@
 
 class AnnotationSpec(proto.Message):
     r"""A definition of an annotation spec.
+
     Attributes:
         name (str):
             Output only. Resource name of the annotation spec. Form:

google/cloud/automl_v1/types/classification.py

@@ -35,6 +35,7 @@ class ClassificationType(proto.Enum):
 
 class ClassificationAnnotation(proto.Message):
     r"""Contains annotation details specific to classification.
+
     Attributes:
         score (float):
             Output only. A confidence estimate between
@@ -87,6 +88,7 @@ class ClassificationEvaluationMetrics(proto.Message):
 
     class ConfidenceMetricsEntry(proto.Message):
         r"""Metrics for a single confidence threshold.
+
         Attributes:
             confidence_threshold (float):
                 Output only. Metrics are computed with an
@@ -162,6 +164,7 @@ class ConfidenceMetricsEntry(proto.Message):
 
     class ConfusionMatrix(proto.Message):
         r"""Confusion matrix of the model running the classification.
+
         Attributes:
             annotation_spec_id (Sequence[str]):
                 Output only. IDs of the annotation specs used in the
@@ -189,6 +192,7 @@ class ConfusionMatrix(proto.Message):
 
         class Row(proto.Message):
             r"""Output only. A row in the confusion matrix.
+
             Attributes:
                 example_count (Sequence[int]):
                     Output only. Value of the specific cell in the confusion

google/cloud/automl_v1/types/data_items.py

@@ -51,6 +51,7 @@ class Image(proto.Message):
 
 class TextSnippet(proto.Message):
     r"""A representation of a text snippet.
+
     Attributes:
         content (str):
             Required. The content of the text snippet as
@@ -74,6 +75,7 @@ class TextSnippet(proto.Message):
 
 class DocumentDimensions(proto.Message):
     r"""Message that describes dimension of a document.
+
     Attributes:
         unit (google.cloud.automl_v1.types.DocumentDimensions.DocumentDimensionUnit):
             Unit of the dimension.
@@ -99,6 +101,7 @@ class DocumentDimensionUnit(proto.Enum):
 
 class Document(proto.Message):
     r"""A structured text document e.g. a PDF.
+
     Attributes:
         input_config (google.cloud.automl_v1.types.DocumentInputConfig):
             An input config specifying the content of the
@@ -181,6 +184,7 @@ class TextSegmentType(proto.Enum):
 
 class ExamplePayload(proto.Message):
     r"""Example data used for training or prediction.
+
     Attributes:
         image (google.cloud.automl_v1.types.Image):
             Example image.

google/cloud/automl_v1/types/detection.py

@@ -30,6 +30,7 @@
 
 class ImageObjectDetectionAnnotation(proto.Message):
     r"""Annotation details for image object detection.
+
     Attributes:
         bounding_box (google.cloud.automl_v1.types.BoundingPoly):
             Output only. The rectangle representing the
@@ -66,6 +67,7 @@ class BoundingBoxMetricsEntry(proto.Message):
 
     class ConfidenceMetricsEntry(proto.Message):
         r"""Metrics for a single confidence threshold.
+
         Attributes:
             confidence_threshold (float):
                 Output only. The confidence threshold value

google/cloud/automl_v1/types/image.py

@@ -33,6 +33,7 @@
 
 class ImageClassificationDatasetMetadata(proto.Message):
     r"""Dataset metadata that is specific to image classification.
+
     Attributes:
         classification_type (google.cloud.automl_v1.types.ClassificationType):
             Required. Type of the classification problem.
@@ -44,11 +45,13 @@ class ImageClassificationDatasetMetadata(proto.Message):
 
 
 class ImageObjectDetectionDatasetMetadata(proto.Message):
-    r"""Dataset metadata specific to image object detection.    """
+    r"""Dataset metadata specific to image object detection.
+    """
 
 
 class ImageClassificationModelMetadata(proto.Message):
     r"""Model metadata for image classification.
+
     Attributes:
         base_model_id (str):
             Optional. The ID of the ``base`` model. If it is specified,
@@ -150,6 +153,7 @@ class ImageClassificationModelMetadata(proto.Message):
 
 class ImageObjectDetectionModelMetadata(proto.Message):
     r"""Model metadata specific to image object detection.
+
     Attributes:
         model_type (str):
             Optional. Type of the model. The available values are:
@@ -232,6 +236,7 @@ class ImageObjectDetectionModelMetadata(proto.Message):
 
 class ImageClassificationModelDeploymentMetadata(proto.Message):
     r"""Model deployment metadata specific to Image Classification.
+
     Attributes:
         node_count (int):
             Input only. The number of nodes to deploy the model on. A
@@ -247,6 +252,7 @@ class ImageClassificationModelDeploymentMetadata(proto.Message):
 
 class ImageObjectDetectionModelDeploymentMetadata(proto.Message):
     r"""Model deployment metadata specific to Image Object Detection.
+
     Attributes:
         node_count (int):
             Input only. The number of nodes to deploy the model on. A

google/cloud/automl_v1/types/io.py

@@ -1446,6 +1446,7 @@ class BatchPredictOutputConfig(proto.Message):
 
 class ModelExportOutputConfig(proto.Message):
     r"""Output configuration for ModelExport Action.
+
     Attributes:
         gcs_destination (google.cloud.automl_v1.types.GcsDestination):
             Required. The Google Cloud Storage location where the model
@@ -1518,6 +1519,7 @@ class ModelExportOutputConfig(proto.Message):
 
 class GcsSource(proto.Message):
     r"""The Google Cloud Storage location for the input content.
+
     Attributes:
         input_uris (Sequence[str]):
             Required. Google Cloud Storage URIs to input files, up to

google/cloud/automl_v1/types/model.py

@@ -26,6 +26,7 @@
 
 class Model(proto.Message):
     r"""API proto representing a trained machine learning model.
+
     Attributes:
         translation_model_metadata (google.cloud.automl_v1.types.TranslationModelMetadata):
             Metadata for translation models.

google/cloud/automl_v1/types/model_evaluation.py

@@ -30,6 +30,7 @@
 
 class ModelEvaluation(proto.Message):
     r"""Evaluation results of a model.
+
     Attributes:
         classification_evaluation_metrics (google.cloud.automl_v1.types.ClassificationEvaluationMetrics):
             Model evaluation metrics for image, text,

google/cloud/automl_v1/types/operations.py

@@ -137,31 +137,38 @@ class OperationMetadata(proto.Message):
 
 
 class DeleteOperationMetadata(proto.Message):
-    r"""Details of operations that perform deletes of any entities.    """
+    r"""Details of operations that perform deletes of any entities.
+    """
 
 
 class DeployModelOperationMetadata(proto.Message):
-    r"""Details of DeployModel operation.    """
+    r"""Details of DeployModel operation.
+    """
 
 
 class UndeployModelOperationMetadata(proto.Message):
-    r"""Details of UndeployModel operation.    """
+    r"""Details of UndeployModel operation.
+    """
 
 
 class CreateDatasetOperationMetadata(proto.Message):
-    r"""Details of CreateDataset operation.    """
+    r"""Details of CreateDataset operation.
+    """
 
 
 class CreateModelOperationMetadata(proto.Message):
-    r"""Details of CreateModel operation.    """
+    r"""Details of CreateModel operation.
+    """
 
 
 class ImportDataOperationMetadata(proto.Message):
-    r"""Details of ImportData operation.    """
+    r"""Details of ImportData operation.
+    """
 
 
 class ExportDataOperationMetadata(proto.Message):
     r"""Details of ExportData operation.
+
     Attributes:
         output_info (google.cloud.automl_v1.types.ExportDataOperationMetadata.ExportDataOutputInfo):
             Output only. Information further describing
@@ -188,6 +195,7 @@ class ExportDataOutputInfo(proto.Message):
 
 class BatchPredictOperationMetadata(proto.Message):
     r"""Details of BatchPredict operation.
+
     Attributes:
         input_config (google.cloud.automl_v1.types.BatchPredictInputConfig):
             Output only. The input config that was given
@@ -221,6 +229,7 @@ class BatchPredictOutputInfo(proto.Message):
 
 class ExportModelOperationMetadata(proto.Message):
     r"""Details of ExportModel operation.
+
     Attributes:
         output_info (google.cloud.automl_v1.types.ExportModelOperationMetadata.ExportModelOutputInfo):
             Output only. Information further describing