Refactor type annotations (#1197)

* Move all imports for typing into a central file.
   This will make upgrading easier when bumping the minimum Python version in
   the future. It also enables a more consistent code style.
 * Upgrade typing imports to be in line with the new Python 3.10 base-line.
   E.g., import `collections.abc.Mapping` instead of `typing.Mapping`.
 * Moved some imports out of `if TYPE_CHECKING` blocks to help Sphinx produce
   useful type annotations in the docs again.

Author: robsdedude

benchkit/app.py

@@ -20,7 +20,6 @@
 from contextlib import contextmanager
 from multiprocessing import Semaphore
 
-import typing_extensions as te
 from sanic import Sanic
 from sanic.exceptions import (
     BadRequest,
@@ -43,7 +42,7 @@
     from .workloads import Workload
 
 
-T_App: te.TypeAlias = "Sanic[Config, BenchKitContext]"
+T_App: t.TypeAlias = "Sanic[Config, BenchKitContext]"
 
 
 def create_app() -> T_App:

benchkit/workloads.py

@@ -19,13 +19,16 @@
 import asyncio
 import enum
 import typing as t
+from collections.abc import (
+    Iterable,
+    Iterator,
+    Mapping,
+)
 from dataclasses import dataclass
 
-import typing_extensions as te
-
 
 if t.TYPE_CHECKING:
-    from collections.abc import Iterator
+    import typing_extensions as te
 
     from neo4j import (
         AsyncDriver,
@@ -40,7 +43,7 @@
 ]
 
 
-class Workloads(t.Mapping):
+class Workloads(Mapping):
     def __init__(self) -> None:
         self._workloads: dict[str, Workload] = {}
         self._current_id: int = 0
@@ -517,7 +520,7 @@ class _WorkloadQuery:
 
     @classmethod
     def parse_multiple(cls, queries: t.Any) -> list[te.Self]:
-        if not isinstance(queries, t.Iterable):
+        if not isinstance(queries, Iterable):
             raise TypeError("Workload queries must be a list")
         return [cls.parse(query) for query in queries]
 
@@ -544,7 +547,7 @@ def parse(cls, query: t.Any) -> te.Self:
 @dataclass
 class _WorkloadConfig:
     database: str | None
-    routing: te.Literal["r", "w"]
+    routing: t.Literal["r", "w"]
 
     @classmethod
     def parse(cls, data: t.Any) -> te.Self:
@@ -554,7 +557,7 @@ def parse(cls, data: t.Any) -> te.Self:
             if not isinstance(database, str):
                 raise TypeError("Workload database must be a string")
 
-        routing: te.Literal["r", "w"] = "w"
+        routing: t.Literal["r", "w"] = "w"
         if "routing" in data:
             raw_routing = data["routing"]
             if not isinstance(routing, str):

docs/source/conf.py

@@ -17,7 +17,6 @@
 
 import os
 import sys
-import typing
 
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -123,17 +122,7 @@
 # Don't include type hints in function signatures
 autodoc_typehints = "description"
 
-autodoc_type_aliases = {
-    # The code-base uses `import typing_extensions as te`.
-    # Re-write these to use `typing` instead, as Sphinx always resolves against
-    # the latest version of the `typing` module.
-    # This is a work-around to make Sphinx resolve type hints correctly, even
-    # though we're using `from __future__ import annotations`.
-    "te": typing,
-    # Type alias that's only defined and imported if `typing.TYPE_CHECKING`
-    # is `True`.
-    "_TAuth": "typing.Tuple[typing.Any, typing.Any] | Auth | None",
-}
+autodoc_type_aliases = {}
 
 # -- Options for HTML output ----------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

pyproject.toml

@@ -139,6 +139,7 @@ extend-exclude = [
 
 [tool.ruff.lint]
 preview = true  # to get CPY lints
+typing-modules = ["neo4j._typing"]
 extend-ignore = [
     "RUF002",  # allow ’ (RIGHT SINGLE QUOTATION MARK) to be used as an apostrophe (e.g. "it’s")
 

src/neo4j/__init__.py

@@ -14,8 +14,7 @@
 # limitations under the License.
 
 
-import typing as _t
-
+from . import _typing as _t
 from ._api import (  # noqa: F401 dynamic attributes
     NotificationCategory,
     NotificationDisabledCategory,

src/neo4j/_addressing.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import typing as t
 from contextlib import suppress as _suppress
 from socket import (
     AddressFamily,
@@ -25,17 +24,15 @@
     getservbyname,
 )
 
-
-if t.TYPE_CHECKING:
-    import typing_extensions as te
+from . import _typing as t
 
 
 _T = t.TypeVar("_T")
 
 
 if t.TYPE_CHECKING:
 
-    class _WithPeerName(te.Protocol):
+    class _WithPeerName(t.Protocol):
         def getpeername(self) -> tuple: ...
 
 

src/neo4j/_api.py

@@ -16,21 +16,19 @@
 
 from __future__ import annotations
 
-import typing as t
 from enum import Enum
 from urllib.parse import (
     parse_qs,
     urlparse,
 )
 
-from . import api
+from . import (
+    _typing as t,
+    api,
+)
 from .exceptions import ConfigurationError
 
 
-if t.TYPE_CHECKING:
-    import typing_extensions as te
-
-
 __all__ = [
     "DRIVER_BOLT",
     "DRIVER_NEO4J",
@@ -51,14 +49,14 @@
 ]
 
 
-DRIVER_BOLT: te.Final[str] = "DRIVER_BOLT"
-DRIVER_NEO4J: te.Final[str] = "DRIVER_NEO4J"
+DRIVER_BOLT: t.Final[str] = "DRIVER_BOLT"
+DRIVER_NEO4J: t.Final[str] = "DRIVER_NEO4J"
 
-SECURITY_TYPE_NOT_SECURE: te.Final[str] = "SECURITY_TYPE_NOT_SECURE"
-SECURITY_TYPE_SELF_SIGNED_CERTIFICATE: te.Final[str] = (
+SECURITY_TYPE_NOT_SECURE: t.Final[str] = "SECURITY_TYPE_NOT_SECURE"
+SECURITY_TYPE_SELF_SIGNED_CERTIFICATE: t.Final[str] = (
     "SECURITY_TYPE_SELF_SIGNED_CERTIFICATE"
 )
-SECURITY_TYPE_SECURE: te.Final[str] = "SECURITY_TYPE_SECURE"
+SECURITY_TYPE_SECURE: t.Final[str] = "SECURITY_TYPE_SECURE"
 
 
 def parse_neo4j_uri(uri):
@@ -184,7 +182,7 @@ class NotificationMinimumSeverity(str, Enum):
 if t.TYPE_CHECKING:
     T_NotificationMinimumSeverity = (
         NotificationMinimumSeverity
-        | te.Literal[
+        | t.Literal[
             "OFF",
             "WARNING",
             "INFORMATION",
@@ -331,7 +329,7 @@ class NotificationDisabledClassification(str, Enum):
     T_NotificationDisabledCategory = (
         NotificationDisabledCategory
         | NotificationDisabledClassification
-        | te.Literal[
+        | t.Literal[
             "HINT",
             "UNRECOGNIZED",
             "UNSUPPORTED",
@@ -458,5 +456,5 @@ class TelemetryAPI(int, Enum):
 
 
 if t.TYPE_CHECKING:
-    T_RoutingControl = RoutingControl | te.Literal["r", "w"]
+    T_RoutingControl = RoutingControl | t.Literal["r", "w"]
     __all__.append("T_RoutingControl")

src/neo4j/_async/_debug/_concurrency_check.py

@@ -18,10 +18,10 @@
 
 import inspect
 import traceback
-import typing as t
 from copy import deepcopy
 from functools import wraps
 
+from ... import _typing as t
 from ..._async_compat.concurrency import (
     AsyncLock,
     AsyncRLock,

src/neo4j/_async/auth_management.py

@@ -17,9 +17,9 @@
 from __future__ import annotations
 
 import abc
-import typing as t
 from logging import getLogger
 
+from .. import _typing as t
 from .._async_compat.concurrency import (
     AsyncCooperativeLock,
     AsyncLock,
@@ -32,9 +32,11 @@
     ExpiringAuth,
 )
 
+# ignore TCH001 to make sphinx not completely drop the ball
+from ..api import _TAuth  # noqa: TCH001
+
 
 if t.TYPE_CHECKING:
-    from ..api import _TAuth
     from ..exceptions import Neo4jError
 
 

src/neo4j/_async/bookmark_manager.py

@@ -16,8 +16,7 @@
 
 from __future__ import annotations
 
-import typing as t
-
+from .. import _typing as t
 from .._async_compat.concurrency import AsyncCooperativeLock
 from .._async_compat.util import AsyncUtil
 from ..api import (

src/neo4j/_async/config.py

@@ -16,8 +16,7 @@
 
 from __future__ import annotations
 
-import typing as t
-
+from .. import _typing as t
 from .._async_compat.concurrency import AsyncLock
 from .._conf import (
     Config,

src/neo4j/_async/driver.py

@@ -17,19 +17,9 @@
 from __future__ import annotations
 
 import asyncio
-import typing as t
 from types import NoneType
 
-
-if t.TYPE_CHECKING:
-    import ssl
-    import typing_extensions as te
-
-    from .._api import (
-        T_NotificationDisabledCategory,
-        T_NotificationMinimumSeverity,
-    )
-
+from .. import _typing as t
 from .._addressing import Address
 from .._api import (
     DRIVER_BOLT,
@@ -62,6 +52,7 @@
     unit_of_work,
 )
 from ..api import (
+    _TAuth,
     AsyncBookmarkManager,
     Auth,
     BookmarkManager,
@@ -104,10 +95,11 @@
     import ssl
     from enum import Enum
 
-    import typing_extensions as te
-
-    from .._api import T_RoutingControl
-    from ..api import _TAuth
+    from .._api import (
+        T_NotificationDisabledCategory,
+        T_NotificationMinimumSeverity,
+        T_RoutingControl,
+    )
 
     class _DefaultEnum(Enum):
         default = "default"
@@ -635,7 +627,7 @@ async def close(self) -> None:
     @t.overload
     async def execute_query(
         self,
-        query_: te.LiteralString | Query,
+        query_: t.LiteralString | Query,
         parameters_: dict[str, t.Any] | None = None,
         routing_: T_RoutingControl = RoutingControl.WRITE,
         database_: str | None = None,
@@ -653,7 +645,7 @@ async def execute_query(
     @t.overload
     async def execute_query(
         self,
-        query_: te.LiteralString | Query,
+        query_: t.LiteralString | Query,
         parameters_: dict[str, t.Any] | None = None,
         routing_: T_RoutingControl = RoutingControl.WRITE,
         database_: str | None = None,
@@ -668,7 +660,7 @@ async def execute_query(
 
     async def execute_query(
         self,
-        query_: te.LiteralString | Query,
+        query_: t.LiteralString | Query,
         parameters_: dict[str, t.Any] | None = None,
         routing_: T_RoutingControl = RoutingControl.WRITE,
         database_: str | None = None,
@@ -677,7 +669,7 @@ async def execute_query(
             AsyncBookmarkManager
             | BookmarkManager
             | None
-            | te.Literal[_DefaultEnum.default]
+            | t.Literal[_DefaultEnum.default]
         ) = _default,
         auth_: _TAuth = None,
         result_transformer_: t.Callable[
@@ -1299,7 +1291,7 @@ async def _get_server_info(self, session_config) -> ServerInfo:
 
 async def _work(
     tx: AsyncManagedTransaction,
-    query: te.LiteralString,
+    query: t.LiteralString,
     parameters: dict[str, t.Any],
     transformer: t.Callable[[AsyncResult], t.Awaitable[_T]],
 ) -> _T:

src/neo4j/_async/home_db_cache.py

@@ -19,9 +19,9 @@
 from __future__ import annotations
 
 import math
-import typing as t
 from time import monotonic
 
+from .. import _typing as t
 from .._async_compat.concurrency import AsyncCooperativeLock