Investigating if we can drop the typeguard dependency.

diffrax/_typing.py

@@ -1,5 +1,4 @@
 import inspect
-import sys
 import types
 from typing import (
     Annotated,
@@ -14,33 +13,41 @@
 )
 from typing_extensions import TypeAlias
 
-import typeguard
-
 
 # We don't actually care what people have subscripted with.
 # In practice this should be thought of as TypeLike = Union[type, types.UnionType]. Plus
 # maybe type(Literal) and so on?
 TypeLike: TypeAlias = Any
 
 
-def better_isinstance(x, annotation) -> bool:
-    """As `isinstance`, but supports general type hints."""
+_T = TypeVar("_T")
 
-    @typeguard.typechecked
-    def f(y: annotation):
-        pass
 
-    try:
-        f(x)
-    except TypeError:
-        return False
-    else:
-        return True
+class _Foo(Generic[_T]):
+    pass
+
 
+_generic_alias_types = (types.GenericAlias, type(_Foo[int]))
+_union_types = (Union, types.UnionType)
+del _Foo, _T
 
-_union_types: list = [Union]
-if sys.version_info >= (3, 10):
-    _union_types.append(types.UnionType)
+
+def better_isinstance(x, annotation) -> bool:
+    """As `isinstance`, but supports a few other types that are useful to us."""
+    if isinstance(annotation, _generic_alias_types):
+        origin = get_origin(annotation)
+        assert origin is not None
+        return better_isinstance(x, origin)
+    elif isinstance(annotation, _union_types):
+        return any(better_isinstance(x, arg) for arg in get_args(annotation))
+    elif annotation is Any:
+        return True
+    elif isinstance(annotation, type):
+        return isinstance(x, annotation)
+    else:
+        raise NotImplementedError(
+            f"Do not know how to check whether `{x}` is an instance of `{annotation}`."
+        )
 
 
 def get_origin_no_specials(x, error_msg: str) -> Optional[type]: