refactor: migrate off pkg_resources

pkg_resources has been deprecated in favour of the builtin `importlib.*` packages. Note that for python3.7 support, we need the `importlib-metadata` backport library; once that version is dropped, the conditional dependency pin and imports can be removed. See the PyPA docs[1] for more info. 1: https://setuptools.pypa.io/en/latest/pkg_resources.html
googleapis · Feb 12, 2024 · b230f56 · b230f56
1 parent f249353
commit b230f56
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 8 deletions.
diff --git a/.coveragerc b/.coveragerc
@@ -11,8 +11,9 @@ exclude_lines =
     pragma: NO COVER
     # Ignore debug-only repr
     def __repr__
-    # Ignore pkg_resources exceptions.
+    # Ignore importlib.metadata exceptions.
     # This is added at the module level as a safeguard for if someone
     # generates the code and tries to run it without pip installing. This
     # makes it virtually impossible to test properly.
-    except pkg_resources.DistributionNotFound
+    except importlib.metadata.PackageNotFoundError
+    except importlib_metadata.PackageNotFoundError
diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py
@@ -19,9 +19,9 @@
 from concurrent import futures
 import inspect
 import logging
-import pkg_resources  # noqa: F401 # Note this is used after copybara replacement
 import os
 import types
+import sys
 from typing import Iterator, List, Optional, Type, TypeVar, Union
 
 from google.api_core import client_options
@@ -43,6 +43,11 @@
     encryption_spec_v1beta1 as gca_encryption_spec_v1beta1,
 )
 
+if sys.version_info >= (3, 8):
+    from importlib import metadata as importlib_metadata
+else:
+    import importlib_metadata
+
 _TVertexAiServiceClientWithOverride = TypeVar(
     "_TVertexAiServiceClientWithOverride",
     bound=utils.VertexAiServiceClientWithOverride,
@@ -447,9 +452,7 @@ def create_client(
         Returns:
             client: Instantiated Vertex AI Service client with optional overrides
         """
-        gapic_version = pkg_resources.get_distribution(
-            "google-cloud-aiplatform",
-        ).version
+        gapic_version = importlib_metadata.version("google-cloud-aiplatform")
 
         if appended_gapic_version:
             gapic_version = f"{gapic_version}+{appended_gapic_version}"

diff --git a/setup.py b/setup.py
@@ -176,6 +176,7 @@
     install_requires=(
         "google-api-core[grpc] >= 1.34.0, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*",
         "google-auth >= 2.14.1, <3.0.0dev",
+        "importlib-metadata >= 6.0.0, <7.0.0; python_version<'3.8'",
         "proto-plus >= 1.22.0, <2.0.0dev",
         "protobuf>=3.19.5,<5.0.0dev,!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",
         "packaging >= 14.3",

diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py
@@ -18,10 +18,10 @@
 import os
 import uuid
 import pytest
+import sys
 import importlib
 
 import pandas as pd
-import pkg_resources
 import re
 
 from datetime import datetime
@@ -43,6 +43,11 @@
 
 from tests.system.aiplatform import e2e_base
 
+if sys.version_info >= (3, 8):
+    from importlib import metadata
+else:
+    import importlib_metadata as metadata
+
 _TEST_PROJECT = e2e_base._PROJECT
 _TEST_LOCATION = e2e_base._LOCATION
 TEST_BUCKET = os.environ.get(
@@ -302,7 +307,7 @@ def test_create_tabular_dataset_from_dataframe(self, bigquery_dataset):
                 == bigquery.SchemaField(name="datetime_col", field_type="DATETIME")
                 if re.match(
                     r"3.*",
-                    pkg_resources.get_distribution("google-cloud-bigquery").version,
+                    metadata.version("google-cloud-bigquery"),
                 )
                 else bigquery.SchemaField(name="datetime_col", field_type="TIMESTAMP")
             )