chore: fix remaining strict mypy errors

sararob · copybara-github · commit a65a8ef4778f · 2025-04-21T07:04:19.000-07:00
PiperOrigin-RevId: 749786746
diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml
@@ -32,4 +32,4 @@ jobs:
           pip install -r requirements.txt
 
       - name: Run mypy  ${{ matrix.python-version }}
-        run: mypy google/genai/ --config-file=google/genai/mypy.ini
+        run: mypy google/genai/ --strict --config-file=google/genai/mypy.ini
diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py
@@ -20,6 +20,7 @@
 """
 
 import asyncio
+from collections.abc import Awaitable, Generator
 import copy
 from dataclasses import dataclass
 import datetime
@@ -201,7 +202,7 @@ def __aiter__(self) -> 'HttpResponse':
     self.segment_iterator = self.async_segments()
     return self
 
-  async def __anext__(self):
+  async def __anext__(self) -> Any:
     try:
       return await self.segment_iterator.__anext__()
     except StopIteration:
@@ -213,7 +214,7 @@ def json(self) -> Any:
       return ''
     return json.loads(self.response_stream[0])
 
-  def segments(self):
+  def segments(self) -> Generator[Any, None, None]:
     if isinstance(self.response_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       for chunk in self.response_stream:
@@ -222,7 +223,7 @@ def segments(self):
       yield from []
     else:
       # Iterator of objects retrieved from the API.
-      for chunk in self.response_stream.iter_lines():
+      for chunk in self.response_stream.iter_lines():  # type: ignore[union-attr]
         if chunk:
           # In streaming mode, the chunk of JSON is prefixed with "data:" which
           # we must strip before parsing.
@@ -256,7 +257,7 @@ async def async_segments(self) -> AsyncIterator[Any]:
       else:
         raise ValueError('Error parsing streaming response.')
 
-  def byte_segments(self):
+  def byte_segments(self) -> Generator[Union[bytes, Any], None, None]:
     if isinstance(self.byte_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       yield from self.byte_stream
@@ -267,7 +268,7 @@ def byte_segments(self):
           'Byte segments are not supported for streaming responses.'
       )
 
-  def _copy_to_dict(self, response_payload: dict[str, object]):
+  def _copy_to_dict(self, response_payload: dict[str, object]) -> None:
     # Cannot pickle 'generator' object.
     delattr(self, 'segment_iterator')
     for attribute in dir(self):
@@ -521,11 +522,11 @@ def _access_token(self) -> str:
         _refresh_auth(self._credentials)
       if not self._credentials.token:
         raise RuntimeError('Could not resolve API token from the environment')
-      return self._credentials.token
+      return self._credentials.token  # type: ignore[no-any-return]
     else:
       raise RuntimeError('Could not resolve API token from the environment')
 
-  async def _async_access_token(self) -> str:
+  async def _async_access_token(self) -> Union[str, Any]:
     """Retrieves the access token for the credentials asynchronously."""
     if not self._credentials:
       async with self._auth_lock:
@@ -675,7 +676,7 @@ def _request(
 
   async def _async_request(
       self, http_request: HttpRequest, stream: bool = False
-  ):
+  ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     if self.vertexai and not self.api_key:
       http_request.headers['Authorization'] = (
@@ -735,7 +736,7 @@ def request(
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Union[BaseResponse, Any]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -753,7 +754,7 @@ def request_streamed(
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Generator[Any, None, None]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -768,7 +769,7 @@ async def async_request(
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ) -> dict[str, object]:
+  ) -> Union[BaseResponse, Any]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
@@ -785,18 +786,18 @@ async def async_request_streamed(
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ):
+  ) -> Any:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
 
     response = await self._async_request(http_request=http_request, stream=True)
 
-    async def async_generator():
+    async def async_generator():  # type: ignore[no-untyped-def]
       async for chunk in response:
         yield chunk
 
-    return async_generator()
+    return async_generator()  # type: ignore[no-untyped-call]
 
   def upload_file(
       self,
@@ -977,7 +978,7 @@ async def async_upload_file(
 
   async def _async_upload_fd(
       self,
-      file: Union[io.IOBase, anyio.AsyncFile],
+      file: Union[io.IOBase, anyio.AsyncFile[Any]],
       upload_url: str,
       upload_size: int,
       *,
@@ -1093,5 +1094,5 @@ async def async_download_file(
   # This method does nothing in the real api client. It is used in the
   # replay_api_client to verify the response from the SDK method matches the
   # recorded response.
-  def _verify_response(self, response_model: _common.BaseModel):
+  def _verify_response(self, response_model: _common.BaseModel) -> None:
     pass
diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py
@@ -218,7 +218,7 @@ def convert_argument_from_function(
 
 
 def invoke_function_from_dict_args(
-    args: Dict[str, Any], function_to_invoke: Callable
+    args: Dict[str, Any], function_to_invoke: Callable[..., Any]
 ) -> Any:
   converted_args = convert_argument_from_function(args, function_to_invoke)
   try:
@@ -232,7 +232,7 @@ def invoke_function_from_dict_args(
 
 
 async def invoke_function_from_dict_args_async(
-    args: Dict[str, Any], function_to_invoke: Callable
+    args: Dict[str, Any], function_to_invoke: Callable[..., Any]
 ) -> Any:
   converted_args = convert_argument_from_function(args, function_to_invoke)
   try:
@@ -280,7 +280,7 @@ def get_function_response_parts(
 
 async def get_function_response_parts_async(
     response: types.GenerateContentResponse,
-    function_map: dict[str, Callable],
+    function_map: dict[str, Callable[..., Any]],
 ) -> list[types.Part]:
   """Returns the function response parts from the response."""
   func_response_parts = []
diff --git a/google/genai/_replay_api_client.py b/google/genai/_replay_api_client.py
@@ -226,7 +226,7 @@ def initialize_replay_session(self, replay_id: str) -> None:
     self._replay_id = replay_id
     self._initialize_replay_session()
 
-  def _get_replay_file_path(self):
+  def _get_replay_file_path(self) -> str:
     return self._generate_file_path_from_replay_id(
         self.replays_directory, self._replay_id
     )
@@ -575,7 +575,7 @@ def download_file(
 
   async def async_download_file(
       self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None
-  ):
+  ) -> Any:
     self._initialize_replay_session_if_not_loaded()
     request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
diff --git a/google/genai/_transformers.py b/google/genai/_transformers.py
@@ -948,7 +948,7 @@ def t_resolve_operation(api_client: _api_client.BaseApiClient, struct: dict[str,
       if total_seconds > LRO_POLLING_TIMEOUT_SECONDS:
         raise RuntimeError(f'Operation {name} timed out.\n{operation}')
       # TODO(b/374433890): Replace with LRO module once it's available.
-      operation = api_client.request(
+      operation = api_client.request(  # type: ignore[assignment]
           http_method='GET', path=name, request_dict={}
       )
       time.sleep(delay_seconds)
diff --git a/google/genai/chats.py b/google/genai/chats.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 
-from collections.abc import AsyncGenerator
+from collections.abc import Iterator
 import sys
 from typing import AsyncIterator, Awaitable, Optional, Union, get_args
 
@@ -283,7 +283,7 @@ def send_message_stream(
       self,
       message: Union[list[PartUnionDict], PartUnionDict],
       config: Optional[GenerateContentConfigOrDict] = None,
-  ):
+  ) -> Iterator[GenerateContentResponse]:
     """Sends the conversation history with the additional message and yields the model's response in chunks.
 
     Args:
@@ -478,7 +478,7 @@ async def async_generator():  # type: ignore[no-untyped-def]
       chunk = None
       async for chunk in await self._modules.generate_content_stream(  # type: ignore[attr-defined]
           model=self._model,
-          contents=self._curated_history + [input_content],
+          contents=self._curated_history + [input_content],  # type: ignore[arg-type]
           config=config if config else self._config,
       ):
         if not _validate_response(chunk):
@@ -489,13 +489,16 @@ async def async_generator():  # type: ignore[no-untyped-def]
           finish_reason = chunk.candidates[0].finish_reason
         yield chunk
 
+      if not output_contents or finish_reason is None:
+        is_valid = False
+
       self.record_history(
           user_input=input_content,
           model_output=output_contents,
           automatic_function_calling_history=chunk.automatic_function_calling_history,
-          is_valid=is_valid and output_contents and finish_reason,
+          is_valid=is_valid,
       )
-    return async_generator()
+    return async_generator()  # type: ignore[no-untyped-call, no-any-return]
 
 
 class AsyncChats:
diff --git a/google/genai/live.py b/google/genai/live.py
@@ -41,8 +41,8 @@
 
 
 try:
-  from websockets.asyncio.client import ClientConnection  # type: ignore
-  from websockets.asyncio.client import connect  # type: ignore
+  from websockets.asyncio.client import ClientConnection
+  from websockets.asyncio.client import connect
 except ModuleNotFoundError:
   # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
   from websockets.client import ClientConnection  # type: ignore
@@ -127,7 +127,7 @@ async def send_client_content(
           ]
       ] = None,
       turn_complete: bool = True,
-  ):
+  ) -> None:
     """Send non-realtime, turn based content to the model.
 
     There are two ways to send messages to the live API:
@@ -203,7 +203,7 @@ async def send_client_content(
 
     await self._ws.send(json.dumps({'client_content': client_content_dict}))
 
-  async def send_realtime_input(self, *, media: t.BlobUnion):
+  async def send_realtime_input(self, *, media: t.BlobUnion) -> None:
     """Send realtime media chunks to the model.
 
     Use `send_realtime_input` for realtime audio chunks and video
@@ -267,7 +267,7 @@ async def send_tool_response(
           types.FunctionResponseOrDict,
           Sequence[types.FunctionResponseOrDict],
       ],
-  ):
+  ) -> None:
     """Send a tool response to the session.
 
     Use `send_tool_response` to reply to `LiveServerToolCall` messages
diff --git a/google/genai/models.py b/google/genai/models.py
@@ -5565,7 +5565,7 @@ async def _generate_content_stream(
         'post', path, request_dict, http_options
     )
 
-    async def async_generator():
+    async def async_generator():  # type: ignore[no-untyped-def]
       async for response_dict in response_stream:
 
         if self._api_client.vertexai:
@@ -5584,7 +5584,7 @@ async def async_generator():
         self._api_client._verify_response(return_value)
         yield return_value
 
-    return async_generator()
+    return async_generator()  # type: ignore[no-untyped-call, no-any-return]
 
   async def embed_content(
       self,
@@ -6602,13 +6602,13 @@ async def generate_content_stream(
           model=model, contents=contents, config=config
       )
 
-      async def base_async_generator(model, contents, config):
+      async def base_async_generator(model, contents, config):  # type: ignore[no-untyped-def]
         async for chunk in response:  # type: ignore[attr-defined]
           yield chunk
 
-      return base_async_generator(model, contents, config)
+      return base_async_generator(model, contents, config)  # type: ignore[no-untyped-call, no-any-return]
 
-    async def async_generator(model, contents, config):
+    async def async_generator(model, contents, config):  # type: ignore[no-untyped-def]
       remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
       logger.info(
           f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
@@ -6687,15 +6687,15 @@ async def async_generator(model, contents, config):
         )
         contents = t.t_contents(self._api_client, contents)
         if not automatic_function_calling_history:
-          automatic_function_calling_history.extend(contents)  # type: ignore[arg-type]
+          automatic_function_calling_history.extend(contents)
         if isinstance(contents, list) and func_call_content is not None:
-          contents.append(func_call_content)  # type: ignore[arg-type]
-          contents.append(func_response_content)  # type: ignore[arg-type]
+          contents.append(func_call_content)
+          contents.append(func_response_content)
         if func_call_content is not None:
           automatic_function_calling_history.append(func_call_content)
         automatic_function_calling_history.append(func_response_content)
 
-    return async_generator(model, contents, config)
+    return async_generator(model, contents, config)  # type: ignore[no-untyped-call, no-any-return]
 
   async def edit_image(
       self,
diff --git a/google/genai/mypy.ini b/google/genai/mypy.ini
@@ -1,4 +1,7 @@
 [mypy]
-exclude = tests/.*
+exclude = (tests/|_test_api_client\.py)
 plugins = pydantic.mypy
-disable_error_code = import-not-found, import-untyped
+; we are ignoring 'unused-ignore' because we run mypy on Python 3.9 - 3.13 and
+; some errors in _automatic_function_calling_util.py only apply in 3.10+
+; 'import-not-found' and 'import-untyped' are environment specific
+disable_error_code = import-not-found, import-untyped, unused-ignore

Original file line number	Diff line number	Diff line change
`@@ -948,7 +948,7 @@ def t_resolve_operation(api_client: _api_client.BaseApiClient, struct: dict[str,`
`948`	`948`	`if total_seconds > LRO_POLLING_TIMEOUT_SECONDS:`
`949`	`949`	`raise RuntimeError(f'Operation {name} timed out.\n{operation}')`
`950`	`950`	`# TODO(b/374433890): Replace with LRO module once it's available.`
`951`		`- operation = api_client.request(`
	`951`	`+ operation = api_client.request( # type: ignore[assignment]`
`952`	`952`	`http_method='GET', path=name, request_dict={}`
`953`	`953`	`)`
`954`	`954`	`time.sleep(delay_seconds)`