1 files changed, 99 insertions, 0 deletions
diff --git a/dot_ipython/profile_default/startup/45_datetime.py b/dot_ipython/profile_default/startup/45_datetime.py
new file mode 100644
index 0000000..d8dcf4e
--- /dev/null
+++ b/dot_ipython/profile_default/startup/45_datetime.py
@@ -0,0 +1,99 @@
+"""Datetime utils."""
+
+import datetime
+import zoneinfo
+
+# Add some pytz timezones, if pytz available
+try:
+    import pytz
+except ModuleNotFoundError:
+    try:
+        UTC = zoneinfo.ZoneInfo("UTC")
+        LDN = zoneinfo.ZoneInfo("Europe/London")
+        NYC = zoneinfo.ZoneInfo("America/New_York")
+        GST = zoneinfo.ZoneInfo("Asia/Dubai")
+    except KeyError:
+        UTC = datetime.timezone.utc
+        GST = datetime.timezone(datetime.timedelta(hours=4), "Asia/Dubai")
+else:
+    UTC = pytz.timezone("UTC")
+    LDN = pytz.timezone("Europe/London")
+    NYC = pytz.timezone("America/New_York")
+    GST = pytz.timezone("Asia/Dubai")
+
+
+class T(datetime.date):
+    """Convenient dev date wrapper."""
+
+    def __add__(self, other):
+        if other in (D, H, M, S, B):
+            other = other(1)
+        if isinstance(other, int):
+            other = datetime.timedelta(days=other)
+        v = super().__add__(other)
+        if isinstance(v, datetime.date) and not isinstance(
+            v, datetime.datetime
+        ):
+            return type(self)(v.year, v.month, v.day)
+        return v
+
+    def __sub__(self, other):
+        if other in (D, H, M, S, B):
+            other = other(1)
+        if isinstance(other, int):
+            other = datetime.timedelta(days=other)
+        v = super().__sub__(other)
+        if isinstance(v, datetime.date) and not isinstance(
+            v, datetime.datetime
+        ):
+            return type(self)(v.year, v.month, v.day)
+        return v
+
+    def __call__(self, *args):
+        return type(self)(*args)
+
+
+_T = datetime.date.today()
+T = T(_T.year, _T.month, _T.day)
+D = lambda d: datetime.timedelta(days=d)
+H = lambda h: datetime.timedelta(hours=h)
+M = lambda m: datetime.timedelta(minutes=m)
+S = lambda s: datetime.timedelta(seconds=s)
+B = lambda s: BDelta(s)
+
+
+def now(tz=UTC):
+    return datetime.datetime.now(tz=tz)
+
+
+class BDelta:
+    """Convenient dev weekday delta."""
+
+    def __init__(self, n: int):
+        self.n = n
+
+    def __add__(self, other):
+        n = self.n
+        if isinstance(other, datetime.date):
+            if n >= 0:
+                op = lambda a, b: a + b
+            else:
+                op = lambda a, b: a - b
+            while n:
+                while other.weekday() in (5, 6):
+                    other = op(other, datetime.timedelta(days=1))
+                other = op(other, datetime.timedelta(days=1))
+                n = op(n, -1)
+            while other.weekday() in (5, 6):
+                other = op(other, datetime.timedelta(days=1))
+            return T(other.year, other.month, other.day)
+        return NotImplemented
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __rsub__(self, other):
+        return other + B(-self.n)
+
+
+_summarise_startup()