ENH: add strptime_micr helper which would support optional microseconds ([.%f])

Author: yarikoptic

heudiconv/tests/test_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 import json
 from json.decoder import JSONDecodeError
 import os
@@ -21,6 +22,7 @@
     remove_prefix,
     remove_suffix,
     save_json,
+    strptime_micr,
     update_json,
 )
 
@@ -168,6 +170,24 @@ def test_get_datetime() -> None:
     )
 
 
+@pytest.mark.parametrize(
+    "dt, fmt",
+    [
+        ("20230310190100", "%Y%m%d%H%M%S"),
+        ("2023-04-02T11:47:09", "%Y-%m-%dT%H:%M:%S"),
+    ],
+)
+def test_strptime_micr(dt: str, fmt: str) -> None:
+    target = datetime.strptime(dt, fmt)
+    assert strptime_micr(dt, fmt) == target
+    assert strptime_micr(dt, fmt + "[.%f]") == target
+    assert strptime_micr(dt + ".0", fmt + "[.%f]") == target
+    assert strptime_micr(dt + ".000000", fmt + "[.%f]") == target
+    assert strptime_micr(dt + ".1", fmt + "[.%f]") == datetime.strptime(
+        dt + ".1", fmt + ".%f"
+    )
+
+
 def test_remove_suffix() -> None:
     """
     Test utils.remove_suffix()

heudiconv/utils.py

@@ -666,6 +666,27 @@ def get_datetime(date: str, time: str, *, microseconds: bool = True) -> str:
     return datetime_str
 
 
+def strptime_micr(date_string: str, fmt: str) -> datetime:
+    r"""
+    Decorate strptime while supporting optional [.%f] in the format at the end
+
+    Parameters
+    ----------
+    date_string: str
+      Date string to parse
+    fmt: str
+      Format string. If it ends with [.%f], we keep it if date_string ends with
+      '.\d+' regex and not if it does not.
+    """
+
+    optional_micr = "[.%f]"
+    if fmt.endswith(optional_micr):
+        fmt = fmt[: -len(optional_micr)]
+        if re.search(r"\.\d+$", date_string):
+            fmt += ".%f"
+    return datetime.strptime(date_string, fmt)
+
+
 def remove_suffix(s: str, suf: str) -> str:
     """
     Remove suffix from the end of the string