Add support for polars dataframes

gw-moore · Jan 20, 2025 · d451081 · d451081
1 parent 597360a
commit d451081
Show file tree

Hide file tree

Showing 71 changed files with 17,245 additions and 571,209 deletions.
diff --git a/README.md b/README.md
@@ -18,15 +18,26 @@
 
 ## Documentation
 
-The [documentation](https://pyfredapi.readthedocs.io/en/latest/) is made with [Sphinx](https://www.sphinx-doc.org/en/master/) and hosted on [Read the Docs](https://readthedocs.org/).
+The [documentation](https://pyfredapi.readthedocs.io/en/latest/) is made with [MkDocs](https://www.mkdocs.org/) and hosted on [Read the Docs](https://readthedocs.org/).
 
 ## Installation
 
+Install the latest version with pip:
+
 ```bash
 pip install pyfredapi
+```
+
+Install pyfredapi with all optional dependencies.
+
+```bash
+pip install 'pyfredapi[all]'
+```
 
-# install with plotting dependencies
-pip install 'pyfredapi[plot]'
+You can also install a subset of all optional dependencies.
+
+```bash
+pip install 'pyfredapi[polars]'
 ```
 
 ## Quick Start

diff --git a/docs/references/CHANGELOG.md b/docs/references/CHANGELOG.md
@@ -1,5 +1,15 @@
 # Changelog
 
+## Version 0.9.2 - 2024-11-03
+
+### Added
+
+- Support for [polars](https://pola.rs/) as a return format.
+
+### Changed
+
+- Updated optional dependencies.
+
 ## Version 0.9.1 - 2024-11-02
 
 ### Removed

diff --git a/docs/tutorials/category.ipynb b/docs/tutorials/category.ipynb
diff --git a/docs/tutorials/series_collection.ipynb b/docs/tutorials/series_collection.ipynb
diff --git a/pyfredapi/__about__.py b/pyfredapi/__about__.py
@@ -1,3 +1,3 @@
 """Define the package metadata."""
 
-__version__ = "0.9.1"
+__version__ = "0.9.2.dev0"
diff --git a/pyfredapi/category.py b/pyfredapi/category.py
@@ -1,21 +1,24 @@
 """The category module provides functions to request data from the [FRED API Categories endpoints](https://fred.stlouisfed.org/docs/api/fred/#Categories)."""
 
+from __future__ import annotations
+
 from typing import Dict, Literal, Optional
 
 from pydantic import BaseModel, ConfigDict, PositiveInt
 
 from ._base import _get_request
 from .series import SeriesInfo
 from .utils import _convert_pydantic_model_to_frozenset
+from .utils._convert_to_df import _convert_to_pandas, _convert_to_polars
+from .utils.enums import ReturnFormat
+
 from .utils._common_type_hints import (
     ApiKeyType,
-    JsonOrPdType,
+    ReturnTypes,
     JsonType,
     KwargsType,
-    ReturnFmtType,
+    ReturnFormats,
 )
-from .utils._convert_to_pandas import _convert_to_pandas
-from .utils.enums import ReturnFormat
 
 
 class CategoryApiParameters(BaseModel):
@@ -190,9 +193,9 @@ def get_category_series(
 def get_category_tags(
     category_id: Optional[int] = None,
     api_key: ApiKeyType = None,
-    return_format: ReturnFmtType = "json",
+    return_format: ReturnFormats = "json",
     **kwargs: KwargsType,
-) -> JsonOrPdType:
+) -> ReturnTypes:
     """Get the FRED tags for a category by category ID.  [Endpoint documentation](https://fred.stlouisfed.org/docs/api/fred/category_tags.html).
 
     Parameters
@@ -208,8 +211,7 @@ def get_category_tags(
 
     Returns
     -------
-    dict | pd.DataFrame
-        Either a dictionary representing the json response or a pandas dataframe.
+    dict | pd.DataFrame | pl.DataFrame
 
     """
     return_format = ReturnFormat(return_format)
@@ -225,15 +227,17 @@ def get_category_tags(
 
     if return_format == ReturnFormat.pandas:
         return _convert_to_pandas(response["tags"])
+    if return_format == ReturnFormat.polars:
+        return _convert_to_polars(response["tags"])
     return response
 
 
 def get_category_related_tags(
     category_id: Optional[int] = None,
     api_key: ApiKeyType = None,
-    return_format: ReturnFmtType = "json",
+    return_format: ReturnFormats = "json",
     **kwargs: KwargsType,
-) -> JsonOrPdType:
+) -> ReturnTypes:
     """Get the related FRED tags for a category by category ID. [Endpoint documentation](https://fred.stlouisfed.org/docs/api/fred/category_related_tags.html).
 
     Parameters
@@ -249,8 +253,7 @@ def get_category_related_tags(
 
     Returns
     -------
-    dict | pd.DataFrame
-        Either a dictionary representing the json response or a pandas dataframe.
+    dict | pd.DataFrame | pl.DataFrame
 
     """
     return_format = ReturnFormat(return_format)
@@ -266,4 +269,6 @@ def get_category_related_tags(
 
     if return_format == ReturnFormat.pandas:
         return _convert_to_pandas(response["tags"])
+    if return_format == ReturnFormat.polars:
+        return _convert_to_polars(response["tags"])
     return response
diff --git a/pyfredapi/maps.py b/pyfredapi/maps.py
@@ -11,7 +11,7 @@
 
 from ._base import _get_request
 from .utils import _convert_pydantic_model_to_frozenset
-from .utils._common_type_hints import ApiKeyType, JsonType, ReturnFmtType
+from .utils._common_type_hints import ApiKeyType, JsonType, ReturnFormats
 from .utils.enums import ReturnFormat
 
 _geo_fred_url = "https://api.stlouisfed.org/geofred/"
@@ -133,7 +133,7 @@ def get_geoseries(
     api_key: ApiKeyType = None,
     start_date: Optional[str] = None,
     end_date: Optional[str] = None,
-    return_format: ReturnFmtType = "pandas",
+    return_format: ReturnFormats = "pandas",
 ) -> GeoseriesData:
     """Request a cross section of regional data for a specified release date. If no date is specified, the most recent data available are returned.