From e015350400d4d17994401c7aff68aa590b6f6ad7 Mon Sep 17 00:00:00 2001
From: Stijn de Gooijer <stijndegooijer@gmail.com>
Date: Mon, 8 Jan 2024 12:54:48 +0100
Subject: [PATCH] feat(python): Add compact syntax for `int_range` starting
 from 0 (#13530)

---
 py-polars/polars/functions/range/int_range.py | 58 ++++++++++++-------
 .../unit/functions/range/test_int_range.py    | 12 ++++
 2 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/py-polars/polars/functions/range/int_range.py b/py-polars/polars/functions/range/int_range.py
index 5928bf435450..d58d45a90a30 100644
--- 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
index 9261b97a37c1..7f6a8b5c23f2 100644
--- 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])