Import cleanup (#259)

* Use typing_extensions.TypedDict, remove unused imports

* Enforce typing_extensions.TypedDict

Fixes: #209

* Add lifetime instructions

Author: MarkDaoust

google/generativeai/types/answer_types.py

@@ -14,17 +14,10 @@
 # limitations under the License.
 from __future__ import annotations
 
-import abc
-import dataclasses
-from typing import Any, Dict, List, TypedDict, Optional, Union
+from typing import Union
 
 import google.ai.generativelanguage as glm
 
-from google.generativeai import string_utils
-from google.generativeai.types import safety_types
-from google.generativeai.types import citation_types
-from google.generativeai.types import content_types
-
 __all__ = ["Answer"]
 
 FinishReason = glm.Candidate.FinishReason

google/generativeai/types/content_types.py

@@ -5,7 +5,8 @@
 import inspect
 import mimetypes
 import typing
-from typing import Any, Callable, TypedDict, Union
+from typing import Any, Callable, Union
+from typing_extensions import TypedDict
 
 import pydantic
 

google/generativeai/types/file_types.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import datetime
-from typing import TypedDict
 
 from google.generativeai.client import get_default_file_client
 

google/generativeai/types/generation_types.py

@@ -21,7 +21,8 @@
 import dataclasses
 import itertools
 import textwrap
-from typing import TypedDict, Union
+from typing import Union
+from typing_extensions import TypedDict
 
 import google.protobuf.json_format
 import google.api_core.exceptions

tests/test_typing_extensions.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pathlib
+import re
+
+from absl.testing import absltest
+
+TYPING_RE = re.compile(r"\btyping\b.*TypedDict")
+
+
+class TypingExtensionsTests(absltest.TestCase):
+    """Pydantic users need the improved version of TypedDict, from typing_extensions.
+
+    This is only required for python versions <3.12. Once 3.12 is the lowest supported version we can drop this.
+
+    ref: https://docs.pydantic.dev/2.3/usage/types/dicts_mapping/
+
+    > the typing-extensions package is required for Python <3.12
+    """
+
+    def test_no_typing_typed_dict(self):
+        root = pathlib.Path(__file__).parent.parent
+        for fpath in root.rglob("*.py"):
+            source = fpath.read_text()
+            if match := TYPING_RE.search(source):
+                raise ValueError(
+                    f"Please import `TypedDict` from `typing_extensions` not `typing`, in:", fpath
+                )
+
+
+if __name__ == "__main__":
+    absltest.main()