summary refs log tree commit diff
path: root/dot_ipython/profile_default/startup/45_datetime.py
blob: d8dcf4e81cd07c1e092a0ee99e2b4c2acdfbade3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
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()