Move Macros to Task SDK

airflow/macros/__init__.py

@@ -17,60 +17,17 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
-from typing import TYPE_CHECKING, Any
-
-import dateutil  # noqa: F401
-from babel import Locale
-from babel.dates import LC_TIME, format_datetime
-
-import airflow.utils.yaml as yaml  # noqa: F401
-from airflow.sdk.definitions.macros import ds_add, ds_format, json, time, uuid  # noqa: F401
-
-if TYPE_CHECKING:
-    from pendulum import DateTime
-
-
-def ds_format_locale(
-    ds: str, input_format: str, output_format: str, locale: Locale | str | None = None
-) -> str:
-    """
-    Output localized datetime string in a given Babel format.
-
-    :param ds: Input string which contains a date.
-    :param input_format: Input string format (e.g., '%Y-%m-%d').
-    :param output_format: Output string Babel format (e.g., `yyyy-MM-dd`).
-    :param locale: Locale used to format the output string (e.g., 'en_US').
-                   If locale not specified, default LC_TIME will be used and if that's also not available,
-                   'en_US' will be used.
-
-    >>> ds_format("2015-01-01", "%Y-%m-%d", "MM-dd-yy")
-    '01-01-15'
-    >>> ds_format("1/5/2015", "%m/%d/%Y", "yyyy-MM-dd")
-    '2015-01-05'
-    >>> ds_format("12/07/2024", "%d/%m/%Y", "EEEE dd MMMM yyyy", "en_US")
-    'Friday 12 July 2024'
-
-    .. versionadded:: 2.10.0
-    """
-    return format_datetime(
-        datetime.strptime(str(ds), input_format),
-        format=output_format,
-        locale=locale or LC_TIME or Locale("en_US"),
-    )
-
-
-# TODO: Task SDK: Move this to the Task SDK once we evaluate "pendulum"'s dependency
-def datetime_diff_for_humans(dt: Any, since: DateTime | None = None) -> str:
-    """
-    Return a human-readable/approximate difference between datetimes.
-
-    When only one datetime is provided, the comparison will be based on now.
-
-    :param dt: The datetime to display the diff for
-    :param since: When to display the date from. If ``None`` then the diff is
-        between ``dt`` and now.
-    """
-    import pendulum
-
-    return pendulum.instance(dt).diff_for_humans(since)
+from airflow.sdk.definitions.macros import (  # noqa: F401
+    datetime,
+    datetime_diff_for_humans,
+    dateutil,
+    ds_add,
+    ds_format,
+    ds_format_locale,
+    json,
+    random,
+    time,
+    timedelta,
+    uuid,
+    yaml,
+)

task_sdk/pyproject.toml

@@ -28,6 +28,9 @@ dependencies = [
     "jinja2>=3.1.4",
     "methodtools>=0.4.7",
     "msgspec>=0.19.0",
+    'pendulum>=2.1.2,<4.0;python_version<"3.12"',
+    'pendulum>=3.0.0,<4.0;python_version>="3.12"',
+    "python-dateutil>=2.7.0",
     "psutil>=6.1.0",
     "structlog>=24.4.0",
     "retryhttp>=1.2.0,!=1.3.0",

task_sdk/src/airflow/sdk/definitions/macros.py

@@ -23,6 +23,15 @@
 import uuid  # noqa: F401
 from datetime import datetime, timedelta
 from random import random  # noqa: F401
+from typing import TYPE_CHECKING, Any
+
+import dateutil  # noqa: F401
+
+import airflow.utils.yaml as yaml  # noqa: F401
+
+if TYPE_CHECKING:
+    from babel import Locale
+    from pendulum import DateTime
 
 
 def ds_add(ds: str, days: int) -> str:
@@ -59,3 +68,50 @@ def ds_format(ds: str, input_format: str, output_format: str) -> str:
     'Friday 12 July 2024'
     """
     return datetime.strptime(str(ds), input_format).strftime(output_format)
+
+
+def datetime_diff_for_humans(dt: Any, since: DateTime | None = None) -> str:
+    """
+    Return a human-readable/approximate difference between datetimes.
+
+    When only one datetime is provided, the comparison will be based on now.
+
+    :param dt: The datetime to display the diff for
+    :param since: When to display the date from. If ``None`` then the diff is
+        between ``dt`` and now.
+    """
+    import pendulum
+
+    return pendulum.instance(dt).diff_for_humans(since)
+
+
+def ds_format_locale(
+    ds: str, input_format: str, output_format: str, locale: Locale | str | None = None
+) -> str:
+    """
+    Output localized datetime string in a given Babel format.
+
+    :param ds: Input string which contains a date.
+    :param input_format: Input string format (e.g., '%Y-%m-%d').
+    :param output_format: Output string Babel format (e.g., `yyyy-MM-dd`).
+    :param locale: Locale used to format the output string (e.g., 'en_US').
+                   If locale not specified, default LC_TIME will be used and if that's also not available,
+                   'en_US' will be used.
+
+    >>> ds_format("2015-01-01", "%Y-%m-%d", "MM-dd-yy")
+    '01-01-15'
+    >>> ds_format("1/5/2015", "%m/%d/%Y", "yyyy-MM-dd")
+    '2015-01-05'
+    >>> ds_format("12/07/2024", "%d/%m/%Y", "EEEE dd MMMM yyyy", "en_US")
+    'Friday 12 July 2024'
+
+    .. versionadded:: 2.10.0
+    """
+    from babel import Locale
+    from babel.dates import LC_TIME, format_datetime
+
+    return format_datetime(
+        datetime.strptime(str(ds), input_format),
+        format=output_format,
+        locale=locale or LC_TIME or Locale("en_US"),
+    )

task_sdk/tests/definitions/test_macros.py

@@ -17,7 +17,10 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import datetime
+
 import lazy_object_proxy
+import pendulum
 import pytest
 
 from airflow.sdk.definitions import macros
@@ -70,3 +73,71 @@ def test_ds_format(ds, input_format, output_format, expected):
 def test_json_loads(input_value, expected):
     result = macros.json.loads(input_value)
     assert result == expected
+
+
+@pytest.mark.parametrize(
+    "ds, input_format, output_format, locale, expected",
+    [
+        ("2015-01-02", "%Y-%m-%d", "MM-dd-yy", None, "01-02-15"),
+        ("2015-01-02", "%Y-%m-%d", "yyyy-MM-dd", None, "2015-01-02"),
+        ("1/5/2015", "%m/%d/%Y", "MM-dd-yy", None, "01-05-15"),
+        ("1/5/2015", "%m/%d/%Y", "yyyy-MM-dd", None, "2015-01-05"),
+        ("12/07/2024", "%d/%m/%Y", "EEEE dd MMMM yyyy", "en_US", "Friday 12 July 2024"),
+        ("12/07/2024", "%d/%m/%Y", "EEEE dd MMMM yyyy", "nl_BE", "vrijdag 12 juli 2024"),
+        (lazy_object_proxy.Proxy(lambda: "2015-01-02"), "%Y-%m-%d", "MM-dd-yy", None, "01-02-15"),
+        (lazy_object_proxy.Proxy(lambda: "2015-01-02"), "%Y-%m-%d", "yyyy-MM-dd", None, "2015-01-02"),
+        (lazy_object_proxy.Proxy(lambda: "1/5/2015"), "%m/%d/%Y", "MM-dd-yy", None, "01-05-15"),
+        (lazy_object_proxy.Proxy(lambda: "1/5/2015"), "%m/%d/%Y", "yyyy-MM-dd", None, "2015-01-05"),
+    ],
+)
+def test_ds_format_locale(ds, input_format, output_format, locale, expected):
+    result = macros.ds_format_locale(ds, input_format, output_format, locale)
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "dt, since, expected",
+    [
+        (
+            pendulum.datetime(2017, 1, 2),
+            None,
+            pendulum.instance(pendulum.datetime(2017, 1, 2)).diff_for_humans(),
+        ),
+        (pendulum.datetime(2017, 1, 2), pendulum.datetime(2017, 1, 3), "1 day before"),
+        (pendulum.datetime(2017, 1, 2), pendulum.datetime(2017, 1, 1), "1 day after"),
+        (
+            lazy_object_proxy.Proxy(lambda: datetime(2017, 1, 2)),
+            None,
+            pendulum.instance(datetime(2017, 1, 2)).diff_for_humans(),
+        ),
+        (
+            lazy_object_proxy.Proxy(lambda: datetime(2017, 1, 2)),
+            pendulum.datetime(2017, 1, 3),
+            "1 day before",
+        ),
+        (
+            lazy_object_proxy.Proxy(lambda: pendulum.datetime(2017, 1, 2)),
+            pendulum.datetime(2017, 1, 1),
+            "1 day after",
+        ),
+    ],
+)
+def test_datetime_diff_for_humans(dt, since, expected):
+    result = macros.datetime_diff_for_humans(dt, since)
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    "input_value, expected",
+    [
+        ('{"field1":"value1", "field2":4, "field3":true}', {"field1": "value1", "field2": 4, "field3": True}),
+        ("field1: value1\nfield2: value2", {"field1": "value1", "field2": "value2"}),
+        (
+            'field1: [ 1, 2, 3, 4, 5 ]\nfield2: {"mini1" : 1, "mini2" : "2"}',
+            {"field1": [1, 2, 3, 4, 5], "field2": {"mini1": 1, "mini2": "2"}},
+        ),
+    ],
+)
+def test_yaml_loads(input_value, expected):
+    result = macros.yaml.safe_load(input_value)
+    assert result == expected