feat: Refactor and enhance type hints across multiple modules

rath/links/aiohttp.py

@@ -2,7 +2,7 @@
 from http import HTTPStatus
 import json
 from ssl import SSLContext
-from typing import Any, Dict, List, Optional, Type, AsyncIterator
+from typing import Any, Dict, List, Optional, Self, Type, AsyncIterator
 from rath.links.types import Payload
 import aiohttp
 from graphql import OperationType
@@ -43,7 +43,11 @@ class AIOHttpLink(AsyncTerminatingLink):
     """ssl_context is the SSLContext to use for the aiohttp session. By default, this
     is a context that uses the certifi CA bundle."""
 
-    auth_errors: List[HTTPStatus] = Field(default_factory=lambda: (HTTPStatus.FORBIDDEN,))
+    auth_errors: List[HTTPStatus] = Field(
+        default_factory=lambda: [
+            HTTPStatus.FORBIDDEN,
+        ]
+    )
     """auth_errors is a list of HTTPStatus codes that indicate that the request was
     unauthorized. By default, this is just HTTPStatus.FORBIDDEN, but you can
     override this to include other status codes that indicate that the request was
@@ -56,9 +60,9 @@ class AIOHttpLink(AsyncTerminatingLink):
 
     _connected = False
 
-    async def __aenter__(self) -> None:
+    async def __aenter__(self) -> Self:
         """Entery point for the async context manager"""
-        pass
+        return self
 
     async def aconnect(self, operation: Operation) -> None:
         """Connects the link to the server
@@ -69,7 +73,12 @@ async def aconnect(self, operation: Operation) -> None:
         """
         self._connected = True
 
-    async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], traceback: Optional[Any]) -> None:
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        traceback: Optional[Any],
+    ) -> None:
         """Exit point for the async context manager"""
         pass
 

rath/links/base.py

@@ -1,4 +1,4 @@
-from typing import Any, AsyncIterator, Optional, Type
+from typing import Any, AsyncIterator, Optional, Self, Type
 from koil.composition import KoiledModel
 from rath.operation import GraphQLResult, Operation
 from rath.errors import NotComposedError
@@ -23,18 +23,25 @@ async def adisconnect(self) -> None:
         """A coroutine that is called when the link is disconnected."""
         pass
 
-    async def __aenter__(self) -> None:
+    async def __aenter__(self) -> Self:
         """A coroutine that is called when the link is entered."""
-        pass
-
-    async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], traceback: Optional[Any]) -> None:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        traceback: Optional[Any],
+    ) -> None:
         """A coroutine that is called when the link is exited."""
         pass
 
     def aexecute(self, operation: Operation) -> AsyncIterator[GraphQLResult]:
         """A coroutine that takes an operation and returns an AsyncIterator
         of GraphQLResults. This method should be implemented by subclasses."""
-        raise NotImplementedError(f"Please overwrite the asubscribe method in {self.__class__.__name__}")
+        raise NotImplementedError(
+            f"Please overwrite the asubscribe method in {self.__class__.__name__}"
+        )
 
 
 class TerminatingLink(Link):