feat: Update TimeDuration to support years and months (#99)

* Introduces use of `dateutil.relativedelta` to calculate
  year/month diffs in a human-friendly way

changes/99.feature

@@ -0,0 +1 @@
+Update `validators.TimeDuration` class to support years and months

src/ai/backend/common/validators.py

@@ -29,6 +29,7 @@
 import pwd
 
 import dateutil.tz
+from dateutil.relativedelta import relativedelta
 try:
     import jwt
     jwt_available = True
@@ -461,11 +462,33 @@ def check_and_return(self, value: Any) -> datetime.tzinfo:
 
 
 class TimeDuration(t.Trafaret):
-
+    '''
+    Represent the relative difference between two datetime objects,
+    parsed from human-readable time duration expression strings.
+
+    If you specify years or months, it returns an
+    :class:`dateutil.relativedelta.relativedelta` instance
+    which keeps the human-friendly year and month calculation
+    considering leap years and monthly day count differences,
+    instead of simply multiplying 365 days to the value of years.
+    Otherwise, it returns the stdlib's :class:`datetime.timedelta`
+    instance.
+
+    Example:
+    >>> t = datetime(2020, 2, 29)
+    >>> t + check_and_return(years=1)
+    datetime.datetime(2021, 2, 28, 0, 0)
+    >>> t + check_and_return(years=2)
+    datetime.datetime(2022, 2, 28, 0, 0)
+    >>> t + check_and_return(years=3)
+    datetime.datetime(2023, 2, 28, 0, 0)
+    >>> t + check_and_return(years=4)
+    datetime.datetime(2024, 2, 29, 0, 0)  # preserves the same day of month
+    '''
     def __init__(self, *, allow_negative: bool = False) -> None:
         self._allow_negative = allow_negative
 
-    def check_and_return(self, value: Any) -> datetime.timedelta:
+    def check_and_return(self, value: Any) -> Union[datetime.timedelta, relativedelta]:
         if not isinstance(value, str):
             self._failure('value must be string', value=value)
         if len(value) == 0:
@@ -477,6 +500,16 @@ def check_and_return(self, value: Any) -> datetime.timedelta:
                 if not self._allow_negative and t < 0:
                     self._failure('value must be positive', value=value)
                 return datetime.timedelta(seconds=t)
+            elif value[-2:].isalpha():
+                t = int(value[:-2])
+                if not self._allow_negative and t < 0:
+                    self._failure('value must be positive', value=value)
+                if value[-2:] == 'yr':
+                    return relativedelta(years=t)
+                elif value[-2:] == 'mo':
+                    return relativedelta(months=t)
+                else:
+                    self._failure('value is not a known time duration', value=value)
             else:
                 t = float(value[:-1])
                 if not self._allow_negative and t < 0:

tests/test_validators.py

@@ -1,13 +1,12 @@
-from ipaddress import IPv4Address
-
-from datetime import timedelta
+from datetime import datetime, timedelta
 import enum
-import ipaddress
+from ipaddress import IPv4Address, ip_address
 import multidict
 import pickle
 import os
 import pwd
 
+from dateutil.relativedelta import relativedelta
 import pytest
 import trafaret as t
 import yarl
@@ -189,7 +188,7 @@ def test_delimiter_list():
     assert iv.check('xxx') == ['xxx']
     iv = tx.DelimiterSeperatedList(tx.HostPortPair, delimiter=',')
     assert iv.check('127.0.0.1:6379,127.0.0.1:6380') == \
-        [(ipaddress.ip_address('127.0.0.1'), 6379), (ipaddress.ip_address('127.0.0.1'), 6380)]
+        [(ip_address('127.0.0.1'), 6379), (ip_address('127.0.0.1'), 6380)]
 
 
 def test_string_list():
@@ -418,6 +417,7 @@ def test_json_string():
 
 def test_time_duration():
     iv = tx.TimeDuration()
+    date = datetime(2020, 2, 29)
     with pytest.raises(t.DataError):
         iv.check('')
     assert iv.check('1w') == timedelta(weeks=1)
@@ -428,6 +428,9 @@ def test_time_duration():
     assert iv.check('1') == timedelta(seconds=1)
     assert iv.check('0.5h') == timedelta(minutes=30)
     assert iv.check('0.001') == timedelta(milliseconds=1)
+    assert iv.check('1yr') == relativedelta(years=1)
+    assert iv.check('1mo') == relativedelta(months=1)
+    assert date + iv.check('4yr') == date + relativedelta(years=4)
     with pytest.raises(t.DataError):
         iv.check('-1')
     with pytest.raises(t.DataError):
@@ -444,6 +447,8 @@ def test_time_duration_negative():
     assert iv.check('0.001') == timedelta(milliseconds=1)
     assert iv.check('-1') == timedelta(seconds=-1)
     assert iv.check('-3d') == timedelta(days=-3)
+    assert iv.check('-1yr') == relativedelta(years=-1)
+    assert iv.check('-1mo') == relativedelta(months=-1)
     with pytest.raises(t.DataError):
         iv.check('-a')
     with pytest.raises(t.DataError):