From b920c27518bbe2069bf6bb3a80494379d88db3b4 Mon Sep 17 00:00:00 2001
From: baggiponte <57922983+baggiponte@users.noreply.github.com>
Date: Thu, 22 Feb 2024 01:21:27 +0100
Subject: [PATCH] refactor(ranges): use appropriate Polars type aliases (#160)

---
 functime/ranges.py | 49 +++++++++++++++++++++-------------------------
 1 file changed, 22 insertions(+), 27 deletions(-)

diff --git a/functime/ranges.py b/functime/ranges.py
index 9e64fc1c..f60d33c7 100644
--- a/functime/ranges.py
+++ b/functime/ranges.py
@@ -1,6 +1,7 @@
 from typing import Optional
 
 import polars as pl
+from polars.type_aliases import TimeUnit
 
 from functime.offsets import _strip_freq_alias
 
@@ -10,7 +11,7 @@ def make_future_ranges(
     cutoffs: pl.DataFrame,
     fh: int,
     freq: str,
-    time_unit: Optional[str] = None,
+    time_unit: Optional[TimeUnit] = None,
 ) -> pl.DataFrame:
     """Return pl.DataFrame with columns entity_col, time_col.
 
@@ -18,34 +19,28 @@ def make_future_ranges(
     """
     entity_col = cutoffs.columns[0]
     if freq.endswith("i"):
-        future_ranges = cutoffs.select(
-            [
-                pl.col(entity_col),
-                pl.int_ranges(
-                    pl.col("low") + 1,
-                    pl.col("low") + fh + 1,
-                    step=int(freq[:-1]),
-                    eager=False,
-                ).alias(time_col),
-            ]
+        return cutoffs.select(
+            pl.col(entity_col),
+            pl.int_ranges(
+                pl.col("low") + 1,
+                pl.col("low") + fh + 1,
+                step=int(freq[:-1]),
+                eager=False,
+            ).alias(time_col),
         )
     else:
-        time_unit = time_unit or "us"
         offset_n, offset_alias = _strip_freq_alias(freq)
         # Make date ranges
-        future_ranges = cutoffs.select(
-            [
-                pl.col(entity_col),
-                pl.date_ranges(
-                    pl.col("low"),
-                    pl.col("low")
-                    .dt.offset_by(f"{fh * offset_n}{offset_alias}")
-                    .alias("high"),
-                    interval=freq,
-                    closed="right",
-                    time_unit=time_unit,
-                    eager=False,
-                ).alias(time_col),
-            ]
+        return cutoffs.select(
+            pl.col(entity_col),
+            pl.date_ranges(
+                pl.col("low"),
+                pl.col("low")
+                .dt.offset_by(f"{fh * offset_n}{offset_alias}")
+                .alias("high"),
+                interval=freq,
+                closed="right",
+                time_unit=time_unit or "us",
+                eager=False,
+            ).alias(time_col),
         )
-    return future_ranges