extract param parsing (#1147)

Author: AndreyNikiforov

src/icloudpd/base.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python
 """Main script that uses Click to parse command-line arguments"""
 
-import re
 from multiprocessing import freeze_support
 
 import foundation
@@ -51,7 +50,7 @@
 from icloudpd.paths import clean_filename, local_download_path, remove_unicode_chars
 from icloudpd.server import serve_app
 from icloudpd.status import Status, StatusExchange
-from icloudpd.string_helpers import truncate_middle
+from icloudpd.string_helpers import parse_timestamp_or_timedelta, truncate_middle
 from icloudpd.xmp_sidecar import generate_xmp_file
 from pyicloud_ipd.base import PyiCloudService
 from pyicloud_ipd.exceptions import PyiCloudAPIResponseException
@@ -267,21 +266,20 @@ def skip_created_before_generator(
 ) -> Optional[Union[datetime.datetime, datetime.timedelta]]:
     if formatted is None:
         return None
-    # can be timestamp or timedelta
-    m = re.match(r"(\d+)([dD]{1})", formatted)
-    if m is not None and m.lastindex is not None and m.lastindex == 2:
-        return datetime.timedelta(days=float(m.group(1)))
-
-    # try timestamp
-    try:
-        dt = datetime.datetime.fromisoformat(formatted)
-        if dt.tzinfo is None:
-            dt = dt.astimezone(get_localzone())
-        return dt
-    except Exception as e:
+    result = parse_timestamp_or_timedelta(formatted)
+    if result is None:
         raise ValueError(
-            f"Timestamp {formatted} for --skip-created-before parameter did not parse from ISO format successfully: {e}"
-        ) from e
+            "--skip-created-before parameter did not parse ISO timestamp or interval successfully"
+        )
+    if isinstance(result, datetime.datetime):
+        return ensure_tzinfo(get_localzone(), result)
+    return result
+
+
+def ensure_tzinfo(tz: datetime.tzinfo, input: datetime.datetime) -> datetime.datetime:
+    if input.tzinfo is None:
+        return input.astimezone(tz)
+    return input
 
 
 def locale_setter(_ctx: click.Context, _param: click.Parameter, use_os_locale: bool) -> bool:

src/icloudpd/string_helpers.py

@@ -1,5 +1,9 @@
 """String helper functions"""
 
+import datetime
+import re
+from typing import Optional, Union
+
 
 def truncate_middle(string: str, length: int) -> str:
     """Truncates a string to a maximum length, inserting "..." in the middle"""
@@ -13,3 +17,31 @@ def truncate_middle(string: str, length: int) -> str:
     start_length = length - end_length - 4
     end_length = max(end_length, 1)
     return f"{string[:start_length]}...{string[-end_length:]}"
+
+
+def parse_timedelta(
+    formatted: str,
+) -> Optional[datetime.timedelta]:
+    m = re.match(r"(\d+)([dD]{1})", formatted)
+    if m is not None and m.lastindex is not None and m.lastindex == 2:
+        return datetime.timedelta(days=float(m.group(1)))
+    return None
+
+
+def parse_timestamp(
+    formatted: str,
+) -> Optional[datetime.datetime]:
+    try:
+        dt = datetime.datetime.fromisoformat(formatted)
+        return dt
+    except Exception:
+        return None
+
+
+def parse_timestamp_or_timedelta(
+    formatted: str,
+) -> Optional[Union[datetime.datetime, datetime.timedelta]]:
+    p1 = parse_timedelta(formatted)
+    if p1 is None:
+        return parse_timestamp(formatted)
+    return p1

tests/test_string_helpers.py

@@ -1,6 +1,12 @@
+import datetime
+import random
+import string
+import sys
 from unittest import TestCase
 
-from icloudpd.string_helpers import truncate_middle
+import pytest
+
+from icloudpd.string_helpers import parse_timestamp_or_timedelta, truncate_middle
 
 
 class TruncateMiddleTestCase(TestCase):
@@ -17,3 +23,78 @@ def test_truncate_middle(self) -> None:
         assert truncate_middle("test_filename.jpg", 0) == ""
         with self.assertRaises(ValueError):
             truncate_middle("test_filename.jpg", -1)
+
+
+class ParseTimestampeOrTimeDeltaTestCase(TestCase):
+    def test_totality(self) -> None:
+        characters = string.ascii_letters + string.digits
+        for _case in range(500):
+            target_length = random.randint(0, 100)
+            target_string = "".join(random.choice(characters) for i in range(target_length))
+            _result = parse_timestamp_or_timedelta(target_string)
+            #  not throwing is okay
+
+    def test_naive(self) -> None:
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.000600") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 600
+        )
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.006") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 6000
+        )
+        assert parse_timestamp_or_timedelta("2025-01-02") == datetime.datetime(
+            2025, 1, 2, 0, 0, 0, 0
+        )
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 11), reason="Requires Python 3.11 or higher for full parsing support"
+    )
+    def test_naive_311plus(self) -> None:
+        #  short
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.0006") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 600
+        )
+
+    def test_aware(self) -> None:
+        assert parse_timestamp_or_timedelta(
+            "2025-01-02T03:04:05.000600+08:00"
+        ) == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 600, datetime.timezone(datetime.timedelta(hours=8))
+        )
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.006+08:00") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 6000, datetime.timezone(datetime.timedelta(hours=8))
+        )
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 11), reason="Requires Python 3.11 or higher for full parsing support"
+    )
+    def test_aware_311plus(self) -> None:
+        #  short
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.0006Z") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 600, datetime.timezone.utc
+        )
+        assert parse_timestamp_or_timedelta("2025-01-02T03:04:05.0006+08:00") == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 600, datetime.timezone(datetime.timedelta(hours=8))
+        )
+
+    def test_timestamp_invalid(self) -> None:
+        assert parse_timestamp_or_timedelta("2025-01") is None
+        assert parse_timestamp_or_timedelta("") is None
+
+    def test_delta_totality(self) -> None:
+        digits = string.digits
+        characters = ["d", "D"]
+        for _case in range(500):
+            target_length = random.randint(1, 5)
+            target_string = "".join(
+                random.choice(digits) for i in range(target_length)
+            ) + random.choice(characters)
+            result = parse_timestamp_or_timedelta(target_string)
+            assert isinstance(result, datetime.timedelta)
+
+    def test_delta(self) -> None:
+        assert parse_timestamp_or_timedelta("1d") == datetime.timedelta(days=1)
+        assert parse_timestamp_or_timedelta("0d") == datetime.timedelta(days=0)
+
+    def test_delta_invalid(self) -> None:
+        assert parse_timestamp_or_timedelta("-1d") is None
+        assert parse_timestamp_or_timedelta("") is None