Import typeddict from typing_extentions (#104)

* import typeddict from typing_extentions

* Add typing-extensions

* Move shim logic to one file.

* typo

* simpilfy

* order

---------

Co-authored-by: Mark Daoust <markdaoust@google.com>

google/generativeai/types/citation_types.py

@@ -13,13 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
+from typing import List
 
-from typing import Optional, List
+from typing_extensions import TypedDict
 
 from google.ai import generativelanguage as glm
 from google.generativeai import string_utils
 
-from typing import TypedDict
 
 __all__ = [
     "CitationMetadataDict",

google/generativeai/types/discuss_types.py

@@ -16,14 +16,16 @@
 
 import abc
 import dataclasses
-from typing import Any, Dict, TypedDict, Union, Iterable, Optional, Tuple, List
+from typing import Any, Dict, Union, Iterable, Optional, Tuple, List
+from typing_extensions import TypedDict
 
 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
 
+
 __all__ = [
     "MessageDict",
     "MessageOptions",

google/generativeai/types/model_types.py

@@ -22,12 +22,16 @@
 import json
 import pathlib
 import re
-from typing import Any, Iterable, TypedDict, Union
+
+from typing import Any, Iterable, Union
+
 import urllib.request
+from typing_extensions import TypedDict
 
 import google.ai.generativelanguage as glm
 from google.generativeai import string_utils
 
+
 __all__ = [
     "Model",
     "ModelNameOptions",
@@ -266,12 +270,12 @@ def _convert_dict(data, input_key, output_key):
 
     try:
         inputs = data[input_key]
-    except KeyError as e:
+    except KeyError:
         raise KeyError(f'input_key is "{input_key}", but data has keys: {sorted(data.keys())}')
 
     try:
         outputs = data[output_key]
-    except KeyError as e:
+    except KeyError:
         raise KeyError(f'output_key is "{output_key}", but data has keys: {sorted(data.keys())}')
 
     for i, o in zip(inputs, outputs):

google/generativeai/types/safety_types.py

@@ -2,11 +2,15 @@
 
 from collections.abc import Mapping
 
+import typing
+from typing import Dict, Iterable, List, Union
+
+from typing_extensions import TypedDict
+
+
 from google.ai import generativelanguage as glm
 from google.generativeai import string_utils
 
-import typing
-from typing import Iterable, Dict, Iterable, List, TypedDict, Union
 
 __all__ = [
     "HarmCategory",

google/generativeai/types/text_types.py

@@ -14,9 +14,11 @@
 # limitations under the License.
 from __future__ import annotations
 
+import sys
 import abc
 import dataclasses
-from typing import Any, Dict, List, TypedDict
+from typing import Any, Dict, List
+from typing_extensions import TypedDict
 
 from google.generativeai import string_utils
 from google.generativeai.types import safety_types

setup.py

@@ -46,6 +46,7 @@ def get_version():
     "google-ai-generativelanguage==0.4.0",
     "google-auth",
     "google-api-core",
+    "typing-extensions",
     "protobuf",
     "tqdm",
 ]