feat(python): Add compact syntax for int_range starting from 0 (#13530

)
pola-rs · Jan 8, 2024 · e015350 · e015350
1 parent ac5c6af
commit e015350
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 21 deletions.
diff --git a/py-polars/polars/functions/range/int_range.py b/py-polars/polars/functions/range/int_range.py
@@ -20,8 +20,8 @@
 
 @overload
 def arange(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -32,8 +32,8 @@ def arange(
 
 @overload
 def arange(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -44,8 +44,8 @@ def arange(
 
 @overload
 def arange(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -55,8 +55,8 @@ def arange(
 
 
 def arange(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = 1,
+    end: int | IntoExprColumn | None = None,
     step: int = 1,
     *,
     dtype: PolarsIntegerType = Int64,
@@ -107,8 +107,8 @@ def arange(
 
 @overload
 def int_range(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -119,8 +119,8 @@ def int_range(
 
 @overload
 def int_range(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -131,8 +131,8 @@ def int_range(
 
 @overload
 def int_range(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = ...,
+    end: int | IntoExprColumn | None = ...,
     step: int = ...,
     *,
     dtype: PolarsIntegerType = ...,
@@ -142,8 +142,8 @@ def int_range(
 
 
 def int_range(
-    start: int | IntoExprColumn,
-    end: int | IntoExprColumn,
+    start: int | IntoExprColumn = 0,
+    end: int | IntoExprColumn | None = None,
     step: int = 1,
     *,
     dtype: PolarsIntegerType = Int64,
@@ -155,9 +155,10 @@ def int_range(
     Parameters
     ----------
     start
-        Lower bound of the range (inclusive).
+        Start of the range (inclusive). Defaults to 0.
     end
-        Upper bound of the range (exclusive).
+        End of the range (exclusive). If set to `None` (default),
+        the value of `start` is used and `start` is set to `0`.
     step
         Step size of the range.
     dtype
@@ -186,12 +187,23 @@ def int_range(
             2
     ]
 
+    `end` can be omitted for a shorter syntax.
+
+    >>> pl.int_range(3, eager=True).alias("int")
+    shape: (3,)
+    Series: 'int' [i64]
+    [
+            0
+            1
+            2
+    ]
+
     `int_range` can be used in conjunction with `count` to generate an index column
     for a DataFrame.
 
     >>> df = pl.DataFrame({"a": [1, 3, 5], "b": [2, 4, 6]})
     >>> df.select(
-    ...     pl.int_range(0, pl.count(), dtype=pl.UInt32).alias("index"),
+    ...     pl.int_range(pl.count(), dtype=pl.UInt32).alias("index"),
     ...     pl.all(),
     ... )
     shape: (3, 3)
@@ -205,6 +217,10 @@ def int_range(
     │ 2     ┆ 5   ┆ 6   │
     └───────┴─────┴─────┘
     """
+    if end is None:
+        end = start
+        start = 0
+
     start = parse_as_expression(start)
     end = parse_as_expression(end)
     result = wrap_expr(plr.int_range(start, end, step, dtype))
@@ -265,9 +281,9 @@ def int_ranges(
     Parameters
     ----------
     start
-        Lower bound of the range (inclusive).
+        Start of the range (inclusive).
     end
-        Upper bound of the range (exclusive).
+        End of the range (exclusive).
     step
         Step size of the range.
     dtype

diff --git a/py-polars/tests/unit/functions/range/test_int_range.py b/py-polars/tests/unit/functions/range/test_int_range.py
@@ -37,6 +37,18 @@ def test_int_range() -> None:
     assert_series_equal(pl.select(int_range=result).to_series(), expected)
 
 
+def test_int_range_short_syntax() -> None:
+    result = pl.int_range(3)
+    expected = pl.Series("int", [0, 1, 2])
+    assert_series_equal(pl.select(int=result).to_series(), expected)
+
+
+def test_int_range_start_default() -> None:
+    result = pl.int_range(end=3)
+    expected = pl.Series("int", [0, 1, 2])
+    assert_series_equal(pl.select(int=result).to_series(), expected)
+
+
 def test_int_range_eager() -> None:
     result = pl.int_range(0, 3, eager=True)
     expected = pl.Series("literal", [0, 1, 2])