chore: Fixed all linting issues and enabled linting in vertexai

PiperOrigin-RevId: 638919820

Author: Ark-kun

noxfile.py

@@ -95,7 +95,7 @@ def lint(session):
  "--diff",
  *LINT_PATHS,
  )
- session.run("flake8", "google", "tests")
+ session.run("flake8", *LINT_PATHS)
 
 
 @nox.session(python=DEFAULT_PYTHON_VERSION)
@@ -225,7 +225,7 @@ def unit_ray(session, ray):
 def unit_langchain(session):
  # Install all test dependencies, then install this package in-place.
 
- constraints_path = str(CURRENT_DIRECTORY / "testing" / f"constraints-langchain.txt")
+ constraints_path = str(CURRENT_DIRECTORY / "testing" / "constraints-langchain.txt")
  standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES
  session.install(*standard_deps, "-c", constraints_path)
 
@@ -236,7 +236,7 @@ def unit_langchain(session):
  session.run(
  "py.test",
  "--quiet",
- f"--junitxml=unit_langchain_sponge_log.xml",
+ "--junitxml=unit_langchain_sponge_log.xml",
  "--cov=google",
  "--cov-append",
  "--cov-config=.coveragerc",

tests/unit/vertexai/test_extensions.py

@@ -176,11 +176,7 @@ def execute_extension_mock():
  with mock.patch.object(
  extension_execution_service.ExtensionExecutionServiceClient, "execute_extension"
  ) as execute_extension_mock:
- response_mock = mock.MagicMock()
- response_mock.content.return_value = _TEST_RESPONSE_CONTENT
- api_client_mock = mock.MagicMock()
- api_client_mock.execute_extension.return_value = response_mock
- execute_extension_mock.return_value = api_client_mock
+ execute_extension_mock.return_value.content = _TEST_RESPONSE_CONTENT
  yield execute_extension_mock
 
 

vertexai/_model_garden/_model_garden_models.py

@@ -25,9 +25,6 @@
 from google.cloud.aiplatform import models as aiplatform_models
 from google.cloud.aiplatform import _publisher_models
 
-# this is needed for class registration to _SUBCLASSES
-import vertexai # pylint:disable=unused-import
-
 _SUPPORTED_PUBLISHERS = ["google"]
 
 _SHORT_MODEL_ID_TO_TUNING_PIPELINE_MAP = {

vertexai/extensions/_extensions.py

@@ -302,8 +302,8 @@ def _try_parse_execution_response(
 ) -> Union[_utils.JsonDict, str]:
  content: str = response.content
  try:
- content = json.loads(response.content)
- except:
+ content = json.loads(content)
+ except json.JSONDecodeError:
  pass
  return content
 

vertexai/language_models/_language_models.py

@@ -4163,7 +4163,7 @@ def _uri_join(uri: str, path_fragment: str) -> str:
 
 # Importing here to prevent issues caused by circular references
 # pylint: disable=g-import-not-at-top,g-bad-import-order
-from vertexai.language_models import _distillation
+from vertexai.language_models import _distillation # noqa: E402
 
 
 class _PreviewTextGenerationModel(

vertexai/preview/language_models.py

@@ -23,9 +23,6 @@
  _PreviewTextEmbeddingModel,
  _PreviewTextGenerationModel,
  ChatMessage,
- ChatModel,
- ChatSession,
- CodeChatSession,
  CountTokensResponse,
  InputOutputTextPair,
  TextEmbedding,

vertexai/reasoning_engines/_reasoning_engines.py

@@ -22,6 +22,7 @@
 import typing
 from typing import Optional, Protocol, Sequence, Union
 
+from google.api_core import exceptions
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils as aip_utils
@@ -378,18 +379,14 @@ def _prepare(
  use for staging the artifacts needed.
  extra_packages (Sequence[str]): The set of extra user-provided packages.
  """
- try:
- from google.cloud.exceptions import NotFound
- except:
- NotFound = Exception
  storage = _utils._import_cloud_storage_or_raise()
  cloudpickle = _utils._import_cloudpickle_or_raise()
  storage_client = storage.Client(project=project)
  staging_bucket = staging_bucket.replace("gs://", "")
  try:
  gcs_bucket = storage_client.get_bucket(staging_bucket)
  _LOGGER.info(f"Using bucket {staging_bucket}")
- except NotFound:
+ except exceptions.NotFound:
  new_bucket = storage_client.bucket(staging_bucket)
  gcs_bucket = storage_client.create_bucket(new_bucket, location=location)
  _LOGGER.info(f"Creating bucket {staging_bucket} in {location=}")

vertexai/tuning/_tuning.py

@@ -15,8 +15,7 @@
 # pylint: disable=protected-access
 """Classes to support Tuning."""
 
-import typing
-from typing import Dict, List, Optional, Union, Sequence
+from typing import Dict, List, Optional, Union
 
 from google.auth import credentials as auth_credentials