support lists in assign_columns (#94)

Author: jmoralez

.pre-commit-config.yaml

@@ -0,0 +1,17 @@
+fail_fast: true
+
+repos:
+  - repo: https://github.com/fastai/nbdev
+    rev: 2.2.10
+    hooks:
+      - id: nbdev_clean
+      - id: nbdev_export
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.1
+    hooks:
+      - id: ruff
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.8.0
+    hooks:
+      - id: mypy
+        args: [--ignore-missing-imports]

nbs/processing.ipynb

@@ -202,7 +202,9 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "def assign_columns(df: DataFrame, names: Union[str, List[str]], values: Union[np.ndarray, pd.Series, pl_Series]) -> DataFrame:\n",
+    "def assign_columns(df: DataFrame, names: Union[str, List[str]], values: Union[np.ndarray, pd.Series, pl_Series, List[float]]) -> DataFrame:\n",
+    "    if isinstance(values, list) and (len(values) != df.shape[0] or not isinstance(names, str)):\n",
+    "        raise ValueError('Only single column assignment is supported for lists.')\n",
     "    if isinstance(df, pd.DataFrame):\n",
     "        df[names] = values\n",
     "    else:\n",
@@ -214,9 +216,13 @@
     "            assert isinstance(names, str)\n",
     "            vals = values.alias(names)\n",
     "        else:\n",
-    "            if isinstance(names, str):\n",
-    "                names = [names]\n",
-    "            vals = pl.from_numpy(values, schema=names, orient='row')\n",
+    "            if isinstance(values, np.ndarray):\n",
+    "                if isinstance(names, str):\n",
+    "                    names = [names]\n",
+    "                vals = pl.from_numpy(values, schema=names, orient='row')\n",
+    "            elif isinstance(values, list):\n",
+    "                assert isinstance(names, str)\n",
+    "                vals = pl_Series(name=names, values=values)\n",
     "        df = df.with_columns(vals)\n",
     "    return df"
    ]
@@ -248,12 +254,14 @@
     "    series = assign_columns(series, 'ones', 1)\n",
     "    series = assign_columns(series, 'zeros', np.zeros(series.shape[0]))\n",
     "    series = assign_columns(series, 'as', 'a')\n",
+    "    series = assign_columns(series, 'bs', series.shape[0] * ['b'])\n",
     "    np.testing.assert_allclose(\n",
     "        series[['x', 'y', 'z']],\n",
     "        np.vstack([x, x, x]).T\n",
     "    )\n",
     "    np.testing.assert_equal(series['ones'], np.ones(series.shape[0]))\n",
-    "    np.testing.assert_equal(series['as'], np.full(series.shape[0], 'a'))"
+    "    np.testing.assert_equal(series['as'], np.full(series.shape[0], 'a'))\n",
+    "    np.testing.assert_equal(series['bs'], np.full(series.shape[0], 'b'))"
    ]
   },
   {

utilsforecast/processing.py

@@ -120,8 +120,12 @@ def maybe_compute_sort_indices(
 def assign_columns(
     df: DataFrame,
     names: Union[str, List[str]],
-    values: Union[np.ndarray, pd.Series, pl_Series],
+    values: Union[np.ndarray, pd.Series, pl_Series, List[float]],
 ) -> DataFrame:
+    if isinstance(values, list) and (
+        len(values) != df.shape[0] or not isinstance(names, str)
+    ):
+        raise ValueError("Only single column assignment is supported for lists.")
     if isinstance(df, pd.DataFrame):
         df[names] = values
     else:
@@ -133,9 +137,13 @@ def assign_columns(
             assert isinstance(names, str)
             vals = values.alias(names)
         else:
-            if isinstance(names, str):
-                names = [names]
-            vals = pl.from_numpy(values, schema=names, orient="row")
+            if isinstance(values, np.ndarray):
+                if isinstance(names, str):
+                    names = [names]
+                vals = pl.from_numpy(values, schema=names, orient="row")
+            elif isinstance(values, list):
+                assert isinstance(names, str)
+                vals = pl_Series(name=names, values=values)
         df = df.with_columns(vals)
     return df