Refactor: improve type annotations

tests/test_executor.py

@@ -75,7 +75,7 @@ def test_pagination():
     p_nopag = Path("pag_off", [int])
     p_pag = Path("pag_on", [int], pagination=True)
 
-    res_nopag = executor.run(path=p_nopag)
+    res_nopag = executor.run(path=p_nopag).parsed
     res_pag = executor.run(path=p_pag, page=2, per_page=3)
 
     assert isinstance(res_nopag, list)

trakt/core/executors.py

@@ -18,6 +18,7 @@
 
 from trakt.core import json_parser
 from trakt.core.components.cache import FrozenRequest
+from trakt.core.components.http_component import ApiResponse
 from trakt.core.exceptions import ArgumentError, ClientError
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -26,6 +27,10 @@
     from trakt.core.paths.path import Path
 
 
+T = TypeVar("T")
+PER_PAGE_LIMIT = 100
+
+
 class Executor:
     params: List[str]
     client: TraktApi
@@ -48,7 +53,9 @@ def __repr__(self) -> str:  # pragma: no cover
     def install(self, suites: List[SuiteInterface]) -> None:
         self.path_suites.extend(suites)
 
-    def run(self, *, path: Optional[Path] = None, **kwargs: Any) -> Any:
+    def run(
+        self, *, path: Optional[Path] = None, **kwargs: Any
+    ) -> Union[ApiResponse, PaginationIterator[T]]:
         if not path:
             return self._delegate_to_interface(**kwargs)
 
@@ -59,7 +66,9 @@ def run(self, *, path: Optional[Path] = None, **kwargs: Any) -> Any:
 
         return self.exec_path_call(path, **kwargs)
 
-    def _delegate_to_interface(self, **kwargs):
+    def _delegate_to_interface(
+        self, **kwargs
+    ) -> Union[ApiResponse, PaginationIterator[T]]:
         matching_paths = self.find_matching_path()
 
         if len(matching_paths) != 1:
@@ -70,12 +79,8 @@ def _delegate_to_interface(self, **kwargs):
         return interface_handler(**kwargs)
 
     def exec_path_call(
-        self,
-        path: Path,
-        extra_quargs: Optional[Dict[str, str]] = None,
-        return_extras: bool = False,
-        **kwargs: Any,
-    ):
+        self, path: Path, extra_quargs: Optional[Dict[str, str]] = None, **kwargs: Any
+    ) -> ApiResponse:
         caching_enabled = self._should_use_cache(path, kwargs.get("no_cache", False))
 
         api_path, query_args = path.get_path_and_qargs()
@@ -99,7 +104,7 @@ def exec_path_call(
             last_request = cast(FrozenRequest, self.client.http.last_request)
             self.client.cache.set(last_request)
 
-        return api_response if return_extras else api_response.parsed
+        return api_response
 
     def _should_use_cache(self, path: Path, no_cache: bool):
         return no_cache is False and self.client.cache.accepted_level(path.cache_level)
@@ -115,10 +120,6 @@ def find_matching_path(self) -> List[Tuple[Path, Callable]]:
         return [p for s in self.path_suites for p in s.find_matching(self.params)]
 
 
-T = TypeVar("T")
-PER_PAGE_LIMIT = 100
-
-
 class PaginationIterator(Iterable[T]):
     pages_total: Optional[int] = None
 

trakt/core/paths/suite_interface.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Type, Union
 
+from trakt.core.components.http_component import ApiResponse
 from trakt.core.exceptions import ArgumentError
 from trakt.core.models import Episode, Movie, Person, Season, Show, TraktList
 from trakt.core.paths.response_structs import Comment
@@ -35,8 +36,15 @@ def find_matching(self, name: Union[List[str], str]) -> List[Tuple[Path, Callabl
     def run(self, command: str, **kwargs: Any) -> Any:
         return self.run_path(path=self._get_path(command), **kwargs)
 
-    def run_path(self, path: Path, **kwargs: Any) -> Any:
-        return self.executor_class(self.client).run(path=path, **kwargs)
+    def run_path(self, path: Path, return_extras: bool = False, **kwargs: Any) -> Any:
+        r = self.executor_class(self.client).run(path=path, **kwargs)
+
+        if isinstance(r, ApiResponse):
+            if return_extras:
+                return r
+            return r.parsed
+
+        return r
 
     def _get_path(self, command: str) -> Path:
         return self.paths[command]