fix: use col() references in tests and fix formatting

Author: BABTUNA

daft/functions/datetime.py

@@ -1385,7 +1385,9 @@ def make_timestamp_ltz(
     Examples:
         >>> import daft
         >>> from daft.functions import make_timestamp_ltz
-        >>> make_timestamp_ltz(daft.col("y"), daft.col("m"), daft.col("d"), daft.col("h"), daft.col("mi"), daft.col("s"))
+        >>> make_timestamp_ltz(
+        ...     daft.col("y"), daft.col("m"), daft.col("d"), daft.col("h"), daft.col("mi"), daft.col("s")
+        ... )
         make_timestamp_ltz(col(y), col(m), col(d), col(h), col(mi), col(s))
     """
     year = Expression._to_expression(year)

tests/dataframe/test_temporals.py

@@ -827,14 +827,14 @@ def test_current_temporal_sql() -> None:
 
 def test_make_date() -> None:
     df = daft.from_pydict({"y": [2021, 2020, 2000], "m": [1, 2, 12], "d": [15, 29, 31]})
-    df = df.with_column("dt", make_date("y", "m", "d"))
+    df = df.with_column("dt", make_date(col("y"), col("m"), col("d")))
     result = df.to_pydict()
     assert result["dt"] == [date(2021, 1, 15), date(2020, 2, 29), date(2000, 12, 31)]
 
 
 def test_make_date_invalid() -> None:
     df = daft.from_pydict({"y": [2021, 2021], "m": [2, 13], "d": [30, 1]})
-    df = df.with_column("dt", make_date("y", "m", "d"))
+    df = df.with_column("dt", make_date(col("y"), col("m"), col("d")))
     result = df.to_pydict()
     assert result["dt"] == [None, None]
 
@@ -843,20 +843,16 @@ def test_make_date_invalid() -> None:
 
 
 def test_make_timestamp() -> None:
-    df = daft.from_pydict(
-        {"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [30], "s": [45.0]}
-    )
-    df = df.with_column("ts", make_timestamp("y", "m", "d", "h", "mi", "s"))
+    df = daft.from_pydict({"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [30], "s": [45.0]})
+    df = df.with_column("ts", make_timestamp(col("y"), col("m"), col("d"), col("h"), col("mi"), col("s")))
     result = df.to_pydict()
     assert result["ts"] == [datetime(2021, 1, 1, 12, 30, 45)]
 
 
 def test_make_timestamp_with_timezone() -> None:
-    df = daft.from_pydict(
-        {"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]}
-    )
+    df = daft.from_pydict({"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]})
     df = df.with_column(
-        "ts", make_timestamp("y", "m", "d", "h", "mi", "s", timezone="UTC")
+        "ts", make_timestamp(col("y"), col("m"), col("d"), col("h"), col("mi"), col("s"), timezone="UTC")
     )
     result = df.to_pydict()
     assert result["ts"] == [datetime(2021, 1, 1, 12, 0, 0, tzinfo=timezone.utc)]
@@ -866,20 +862,16 @@ def test_make_timestamp_with_timezone() -> None:
 
 
 def test_make_timestamp_ltz() -> None:
-    df = daft.from_pydict(
-        {"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]}
-    )
-    df = df.with_column("ts", make_timestamp_ltz("y", "m", "d", "h", "mi", "s"))
+    df = daft.from_pydict({"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]})
+    df = df.with_column("ts", make_timestamp_ltz(col("y"), col("m"), col("d"), col("h"), col("mi"), col("s")))
     result = df.to_pydict()
     assert result["ts"] == [datetime(2021, 1, 1, 12, 0, 0, tzinfo=timezone.utc)]
 
 
 def test_make_timestamp_ltz_with_timezone() -> None:
-    df = daft.from_pydict(
-        {"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]}
-    )
+    df = daft.from_pydict({"y": [2021], "m": [1], "d": [1], "h": [12], "mi": [0], "s": [0.0]})
     df = df.with_column(
-        "ts", make_timestamp_ltz("y", "m", "d", "h", "mi", "s", timezone="US/Eastern")
+        "ts", make_timestamp_ltz(col("y"), col("m"), col("d"), col("h"), col("mi"), col("s"), timezone="US/Eastern")
     )
     result = df.to_pydict()
     # 12:00 EST = 17:00 UTC
@@ -890,10 +882,8 @@ def test_make_timestamp_ltz_with_timezone() -> None:
 
 
 def test_last_day() -> None:
-    df = daft.from_pydict(
-        {"dt": [date(2021, 1, 15), date(2021, 2, 10), date(2020, 2, 10), date(2021, 4, 1)]}
-    )
-    df = df.with_column("last", last_day("dt"))
+    df = daft.from_pydict({"dt": [date(2021, 1, 15), date(2021, 2, 10), date(2020, 2, 10), date(2021, 4, 1)]})
+    df = df.with_column("last", last_day(col("dt")))
     result = df.to_pydict()
     assert result["last"] == [
         date(2021, 1, 31),
@@ -911,7 +901,7 @@ def test_next_day() -> None:
     df = daft.from_pydict({"dt": [date(2021, 1, 1), date(2021, 1, 1), date(2021, 1, 1)]})
     results = {}
     for dow in ["Monday", "Friday", "Sunday"]:
-        tmp = df.with_column("nd", next_day("dt", dow))
+        tmp = df.with_column("nd", next_day(col("dt"), dow))
         results[dow] = tmp.to_pydict()["nd"]
     assert results["Monday"] == [date(2021, 1, 4)] * 3  # next Mon
     assert results["Friday"] == [date(2021, 1, 8)] * 3  # next Fri (not same day)