s

Author: afuetterer

src/re3data/_client/_async.py

@@ -22,8 +22,9 @@
 from re3data._response import Response, _build_response
 
 if TYPE_CHECKING:
-    from re3data._resources import Repository, RepositorySummary
+    from pandas import DataFrame
 
+    from re3data._resources import Repository, RepositorySummary
 logger = logging.getLogger(__name__)
 
 
@@ -59,7 +60,7 @@ async def list(
         query: str | None = None,
         return_type: ReturnType = ReturnType.DATACLASS,
         count: bool = False,
-    ) -> list[RepositorySummary] | Response | dict[str, Any] | str | int:
+    ) -> list[RepositorySummary] | Response | dict[str, Any] | DataFrame | str | int:
         """List the metadata of all repositories in the re3data API.
 
         Args:
@@ -83,7 +84,7 @@ async def list(
 
     async def get(
         self, repository_id: str, return_type: ReturnType = ReturnType.DATACLASS
-    ) -> Repository | Response | dict[str, Any] | str:
+    ) -> Repository | Response | dict[str, Any] | DataFrame | str:
         """Get the metadata of a specific repository.
 
         Args:

src/re3data/_client/_sync.py

@@ -24,8 +24,9 @@
 from re3data._response import Response, _build_response
 
 if TYPE_CHECKING:
-    from re3data._resources import Repository, RepositorySummary
+    from pandas import DataFrame
 
+    from re3data._resources import Repository, RepositorySummary
 logger = logging.getLogger(__name__)
 
 
@@ -61,7 +62,7 @@ def list(
         query: str | None = None,
         return_type: ReturnType = ReturnType.DATACLASS,
         count: bool = False,
-    ) -> list[RepositorySummary] | Response | dict[str, Any] | str | int:
+    ) -> list[RepositorySummary] | Response | dict[str, Any] | DataFrame | str | int:
         """List the metadata of all repositories in the re3data API.
 
         Args:
@@ -85,7 +86,7 @@ def list(
 
     def get(
         self, repository_id: str, return_type: ReturnType = ReturnType.DATACLASS
-    ) -> Repository | Response | dict[str, Any] | str:
+    ) -> Repository | Response | dict[str, Any] | DataFrame | str:
         """Get the metadata of a specific repository.
 
         Args:

src/re3data/_client/base.py

@@ -16,6 +16,8 @@
 from re3data._serializer import _to_csv, _to_dataframe, _to_dict, _to_json
 
 if TYPE_CHECKING:
+    from pandas import DataFrame
+
     from re3data._resources import Repository, RepositorySummary
 
 BASE_URL: str = "https://www.re3data.org/api/beta/"
@@ -79,25 +81,22 @@ def _build_query_params(query: str | None = None) -> dict[str, str]:
     return query_params
 
 
-# TODO: add dataframe to return types
-
-
 @overload
 def _dispatch_return_type(
     response: Response, resource_type: Literal[ResourceType.REPOSITORY], return_type: ReturnType, count: bool = False
-) -> Repository | Response | dict[str, Any] | str: ...
+) -> Repository | Response | dict[str, Any] | DataFrame | str: ...
 @overload
 def _dispatch_return_type(
     response: Response,
     resource_type: Literal[ResourceType.REPOSITORY_LIST],
     return_type: ReturnType,
     count: bool = False,
-) -> list[RepositorySummary] | Response | dict[str, Any] | str | int: ...
+) -> list[RepositorySummary] | Response | dict[str, Any] | DataFrame | str | int: ...
 
 
 def _dispatch_return_type(  # noqa: PLR0911
     response: Response, resource_type: ResourceType, return_type: ReturnType, count: bool = False
-) -> Repository | list[RepositorySummary] | Response | dict[str, Any] | str | int:
+) -> Repository | list[RepositorySummary] | Response | dict[str, Any] | DataFrame | str | int:
     """Dispatch the response to the correct return type based on the provided return type and resource type.
 
     Args: