fix timestamp check in time_ranges for timestamps with timezone (#84)

Nixtla · May 15, 2024 · 6d99b22 · 6d99b22
1 parent 9eb7024
commit 6d99b22
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 17 deletions.
diff --git a/nbs/processing.ipynb b/nbs/processing.ipynb
@@ -901,21 +901,25 @@
     "    if isinstance(starts, pd.Series):\n",
     "        starts = pd.Index(starts)\n",
     "    if isinstance(starts, pd.Index):\n",
-    "        starts_dtype = starts.dtype.type\n",
-    "        if issubclass(starts_dtype, np.integer):\n",
+    "        if _is_int_dtype(starts):\n",
+    "            starts_np = starts.to_numpy(copy=False)  # may be pyarrow\n",
     "            out = np.hstack(\n",
-    "                [np.arange(start, start + freq * periods, freq, dtype=starts_dtype) for start in starts]\n",
+    "                [\n",
+    "                    np.arange(start, start + freq * periods, freq, dtype=starts_np.dtype)\n",
+    "                    for start in starts_np\n",
+    "                ]\n",
     "            )\n",
-    "        elif pd.api.types.is_datetime64_dtype(starts_dtype):\n",
+    "        elif _is_dt_dtype(starts):\n",
     "            if isinstance(freq, str):\n",
     "                freq = pd.tseries.frequencies.to_offset(freq)\n",
     "            out = []\n",
     "            for i in range(periods):\n",
     "                out.append([starts + i * freq])\n",
+    "            # pyarrow timestamps don't seem to work with offsets yet, keeping np.vstack\n",
     "            out = np.vstack(out).ravel(order='F')\n",
     "        else:\n",
-    "            raise ValueError(f\"`starts` must be integers or timestamps, got '{starts_dtype}'.\")\n",
-    "        out = pd.Series(out)\n",
+    "            raise ValueError(f\"`starts` must be integers or timestamps, got '{starts.dtype}'.\")\n",
+    "        out = pd.Series(out, dtype=starts.dtype)\n",
     "    else:\n",
     "        try:\n",
     "            is_int = starts.dtype.is_integer()\n",
@@ -962,8 +966,8 @@
     "    pd.Series(pd.to_datetime(['2000-01-01', '2000-03-01', '2010-10-01', '2010-12-01']))\n",
     ")\n",
     "pd.testing.assert_series_equal(\n",
-    "    time_ranges(pd.to_datetime(['2000-01-01', '2010-01-01']), freq=2 * pd.offsets.YearBegin(), periods=2),\n",
-    "    pd.Series(pd.to_datetime(['2000-01-01', '2002-01-01', '2010-01-01', '2012-01-01']))\n",
+    "    time_ranges(pd.to_datetime(['2000-01-01', '2010-01-01']).tz_localize('US/Eastern'), freq=2 * pd.offsets.YearBegin(), periods=2),\n",
+    "    pd.Series(pd.to_datetime(['2000-01-01', '2002-01-01', '2010-01-01', '2012-01-01']).tz_localize('US/Eastern'))\n",
     ")\n",
     "pd.testing.assert_series_equal(\n",
     "    time_ranges(pd.to_datetime(['2000-12-31', '2010-12-31']), freq=2 * pd.offsets.YearEnd(), periods=2),\n",

diff --git a/settings.ini b/settings.ini
@@ -1,7 +1,7 @@
 [DEFAULT]
 repo = utilsforecast
 lib_name = utilsforecast
-version = 0.1.9
+version = 0.1.10
 min_python = 3.8
 license = apache2
 black_formatting = True

diff --git a/utilsforecast/__init__.py b/utilsforecast/__init__.py
@@ -1 +1 @@
-__version__ = "0.1.9"
+__version__ = "0.1.10"
diff --git a/utilsforecast/processing.py b/utilsforecast/processing.py
@@ -413,26 +413,29 @@ def time_ranges(
     if isinstance(starts, pd.Series):
         starts = pd.Index(starts)
     if isinstance(starts, pd.Index):
-        starts_dtype = starts.dtype.type
-        if issubclass(starts_dtype, np.integer):
+        if _is_int_dtype(starts):
+            starts_np = starts.to_numpy(copy=False)  # may be pyarrow
             out = np.hstack(
                 [
-                    np.arange(start, start + freq * periods, freq, dtype=starts_dtype)
-                    for start in starts
+                    np.arange(
+                        start, start + freq * periods, freq, dtype=starts_np.dtype
+                    )
+                    for start in starts_np
                 ]
             )
-        elif pd.api.types.is_datetime64_dtype(starts_dtype):
+        elif _is_dt_dtype(starts):
             if isinstance(freq, str):
                 freq = pd.tseries.frequencies.to_offset(freq)
             out = []
             for i in range(periods):
                 out.append([starts + i * freq])
+            # pyarrow timestamps don't seem to work with offsets yet, keeping np.vstack
             out = np.vstack(out).ravel(order="F")
         else:
             raise ValueError(
-                f"`starts` must be integers or timestamps, got '{starts_dtype}'."
+                f"`starts` must be integers or timestamps, got '{starts.dtype}'."
             )
-        out = pd.Series(out)
+        out = pd.Series(out, dtype=starts.dtype)
     else:
         try:
             is_int = starts.dtype.is_integer()