fix: update abstract type hints to use collections.abc instead of typ…

chex/_src/asserts.py

@@ -19,7 +19,8 @@
 import functools
 import inspect
 import traceback
-from typing import Any, Callable, List, Optional, Sequence, Set, Union, cast
+from collections.abc import Sequence
+from typing import Any, Callable, List, Optional, Set, Union, cast
 import unittest
 from unittest import mock
 

chex/_src/asserts_chexify_test.py

@@ -19,7 +19,8 @@
 import sys
 import threading
 import time
-from typing import Any, Optional, Sequence, Type
+from collections.abc import Sequence
+from typing import Any, Optional, Type
 
 from absl.testing import absltest
 from absl.testing import parameterized

chex/_src/asserts_internal.py

@@ -24,12 +24,12 @@
 
 import collections
 import collections.abc
-from collections.abc import Hashable
+from collections.abc import Hashable, Sequence
 import functools
 import re
 import threading
 import traceback
-from typing import Any, Sequence, Union, Callable, List, Optional, Set, Tuple, Type
+from typing import Any, Union, Callable, List, Optional, Set, Tuple, Type
 
 from absl import logging
 from chex._src import pytypes

chex/_src/dataclass_test.py

@@ -20,7 +20,8 @@
 import dataclasses
 import pickle
 import sys
-from typing import Any, Generic, Mapping, TypeVar
+from collections.abc import Mapping
+from typing import Any, Generic, TypeVar
 import unittest
 
 from absl.testing import absltest

chex/_src/dimensions.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 """Utilities to hold expected dimension sizes."""
 
-from collections.abc import Sized
+from collections.abc import Collection, Sized
 import math
 import re
-from typing import Any, Collection, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 
 Shape = Tuple[Optional[int], ...]

chex/_src/fake.py

@@ -25,7 +25,8 @@
 import inspect
 import os
 import re
-from typing import Any, Callable, Iterable, Optional, Union
+from collections.abc import Iterable
+from typing import Any, Callable, Optional, Union
 from unittest import mock
 from absl import flags
 import jax

chex/_src/pytypes.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Type definitions to use for type annotations."""
 
-from typing import Any, Iterable, Mapping, Sequence, Union
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, Union
 
 import jax
 import numpy as np

chex/_src/restrict_backends.py

@@ -29,7 +29,8 @@
 """
 import contextlib
 import functools
-from typing import Optional, Sequence
+from collections.abc import Sequence
+from typing import Optional
 
 # pylint: disable=g-import-not-at-top
 try:

chex/_src/variants.py

@@ -18,7 +18,8 @@
 import functools
 import inspect
 import itertools
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 import unittest
 
 from absl import flags

chex/_src/warnings.py

@@ -15,7 +15,8 @@
 """Utilities to emit warnings."""
 
 import functools
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 import warnings