A timezone-aware datetime subclass with mixin utilities for easy arithmetic, business-day calculations, human-readable diffs, and seamless pandas integration.

smartdatetime.py

"""
A drop-in extension of Python’s built-in datetime that adds:

Timezone awareness & conversions
Auto-localization, to_utc(), to_local(), to_timezone()

Human-friendly arithmetic
Add/subtract timedelta, seconds, dict-style offsets

Business-day support
add_business_days(), subtract_business_days(), next_weekday()

Boundary helpers
start_of_day(), end_of_day(), start_of_month(), end_of_month()

Date info
week_of_year(), quarter(), is_weekend(), is_weekday()

Serialization & formatting
to_iso(), to_timestamp(), to_dict(), custom __format__ and mapping protocol

Human-readable diffs
humanize_diff() (e.g. “in 2d 3h” or “5h ago”)

Pandas interoperability
Easily convert to pd.Timestamp or datetime64[ns]
"""

import calendar
from datetime import timedelta, datetime
from typing import Any, Dict, Union
import pytz
from tzlocal import get_localzone

class SmartDateTimeMixin:
    """
    Mixin providing additional datetime functionality:
    arithmetic operations, timezone conversions, business day calculations,
    human-readable diffs, and common utilities.
    """
    str_format: str = "%Y-%m-%d %H:%M:%S"
    default_tz = get_localzone()

    def _make_tz_aware(self, dt: datetime) -> datetime:
        """
        Ensure a naive datetime is localized to the default timezone.
        If already timezone-aware, convert it to default_tz.

        :param dt: datetime instance (naive or aware)
        :return: timezone-aware datetime in default_tz
        """
        if dt.tzinfo is None:
            # first try pytz.localize; if not supported, use replace()
            try:
                return self.default_tz.localize(dt)   # pytz path
            except AttributeError:
                return dt.replace(tzinfo=self.default_tz)  # zoneinfo path
        return dt.astimezone(self.default_tz)

    def to_timestamp(self, as_int: bool = True) -> Union[int, float]:
        """
        Return Unix timestamp (seconds since epoch).

        :param as_int: if True, returns int; otherwise, float with fractions.
        :return: timestamp in seconds
        """
        ts = self.timestamp()
        return int(ts) if as_int else ts

    def to_timezone(self, tz_name: str) -> "SmartDateTime":
        """
        Convert this instance to another timezone.

        :param tz_name: IANA timezone string (e.g. 'UTC', 'Europe/London')
        :return: new SmartDateTime in target timezone
        """
        tz = pytz.timezone(tz_name) if tz_name != "UTC" else pytz.UTC
        return self.astimezone(tz)

    def to_utc(self) -> "SmartDateTime":
        """Alias for to_timezone('UTC')."""
        return self.to_timezone("UTC")

    def to_local(self) -> "SmartDateTime":
        """
        Convert to the system's local timezone.
        Requires tzlocal to detect the local timezone.

        :return: new SmartDateTime in local timezone
        """
        local_tz = get_localzone().key
        return self.to_timezone(local_tz)

    def start_of_day(self) -> "SmartDateTime":
        """Return new datetime at 00:00:00 of the same day."""
        return self.replace(hour=0, minute=0, second=0, microsecond=0)

    def end_of_day(self) -> "SmartDateTime":
        """Return new datetime at 23:59:59.999999 of the same day."""
        return self.replace(hour=23, minute=59, second=59, microsecond=999999)

    def is_weekend(self) -> bool:
        """Return True if the date falls on Saturday (6) or Sunday (7)."""
        return self.weekday() >= 5

    def is_weekday(self) -> bool:
        """Return True if the date falls on a weekday (Monday to Friday)."""
        return self.weekday() < 5

    def to_iso(self) -> str:
        """Serialize to complete ISO 8601 string with offset."""
        return self.isoformat()

    def to_dict(self) -> Dict[str, Union[int, str]]:
        """
        Return date/time components as a dictionary:
        year, month, day, hour, minute, second, microsecond, timezone.

        :return: dict of datetime components
        """
        return {
            "year": self.year,
            "month": self.month,
            "day": self.day,
            "hour": self.hour,
            "minute": self.minute,
            "second": self.second,
            "microsecond": self.microsecond,
            "timezone": self.default_tz
        }

    def week_of_year(self) -> int:
        """Return ISO week number (1-53)."""
        return self.isocalendar()[1]

    def quarter(self) -> int:
        """Return quarter of the year (1-4)."""
        return (self.month - 1) // 3 + 1

    def start_of_month(self) -> "SmartDateTime":
        """Return new datetime at the first day of the month, 00:00."""
        return self.replace(day=1, hour=0, minute=0, second=0, microsecond=0)

    def end_of_month(self) -> "SmartDateTime":
        """Return new datetime at the last day of the month, 23:59:59.999999."""
        last_day = calendar.monthrange(self.year, self.month)[1]
        return self.replace(day=last_day, hour=23, minute=59, second=59, microsecond=999999)

    def next_weekday(self, n: int = 1) -> "SmartDateTime":
        """Advance by n business days, skipping weekends."""
        dt = self
        added = 0
        while added < n:
            dt += timedelta(days=1)
            if dt.weekday() < 5:
                added += 1
        return dt

    def add_business_days(self, n: int) -> "SmartDateTime":
        """Alias for next_weekday(n) when n > 0, otherwise self."""
        return self.next_weekday(n) if n > 0 else self

    def subtract_business_days(self, n: int) -> "SmartDateTime":
        """Subtract n business days, skipping weekends."""
        dt = self
        subtracted = 0
        while subtracted < n:
            dt -= timedelta(days=1)
            if dt.weekday() < 5:
                subtracted += 1
        return dt

    def humanize_diff(self, other: Union["SmartDateTime", datetime]) -> str:
        """
        Return a human-readable difference between two dates,
        e.g. "2d 3h ago" or "in 5h".

        :param other: another SmartDateTime or datetime
        :return: human-friendly diff string
        """
        if isinstance(other, SmartDateTimeMixin):
            other_dt = other
        elif isinstance(other, datetime):
            other_dt = other.astimezone(self.default_tz)
        else:
            raise TypeError("other must be SmartDateTime or datetime")
        delta = other_dt - self
        past = delta.total_seconds() < 0
        secs = abs(int(delta.total_seconds()))
        days, rem = divmod(secs, 86400)
        hours, rem = divmod(rem, 3600)
        minutes, seconds = divmod(rem, 60)
        parts = []
        if days:
            parts.append(f"{days}d")
        if hours:
            parts.append(f"{hours}h")
        if minutes:
            parts.append(f"{minutes}m")
        if seconds and not parts:
            parts.append(f"{seconds}s")
        prefix = "in " if not past and parts else ""
        suffix = " ago" if past else ""
        return f"{prefix}{' '.join(parts)}{suffix}"
    

class SmartDateTime(SmartDateTimeMixin, datetime):
    """
    Subclass of datetime integrating SmartDateTimeMixin.
    Supports initialization from None, str, datetime, or numeric timestamp.

    **Usage**:
    SmartDateTime(dt=None|str|datetime|int|float, timezone='UTC', str_format='%Y-%m-%d %H:%M:%S')
    or
    SmartDateTime()  # now() at local timezone
    or
    SmartDateTime(**datetime_kwargs)
    """
    def __new__(cls, *args, timezone=None, str_format=None, **kwargs):
        """
        Create a new SmartDateTime instance.

        :param args/kwargs: either raw datetime components (year, month, ...)
                             or single dt argument (None, str, datetime, number).
        :param timezone: IANA timezone string for naive inputs.
        :param str_format: format string for parsing date strings.
        :return: SmartDateTime instance
        """
        # Raw datetime constructor path
        if len(args) > 1 or len(kwargs) > 1:
            # interpreta como datetime.raw constructor
            obj = datetime.__new__(cls, *args, **kwargs)
            # ajusta fuso (se passado) ou deixa o tzinfo que veio
            tz = pytz.timezone(timezone) if timezone else obj.tzinfo or cls.default_tz
            obj = obj.replace(tzinfo=tz)
            obj.str_format = str_format or cls.str_format
            obj.default_tz = tz
            return obj
        # Single argument dt path
        dt = args[0] if args else kwargs.get('dt', None)

        tz = pytz.timezone(timezone) if timezone else cls.default_tz
        fmt = str_format or cls.str_format

        if dt is None:
            naive = datetime.now()
        elif isinstance(dt, str):
            try:
                naive = datetime.fromisoformat(dt)
            except ValueError:
                naive = datetime.strptime(dt, fmt)
        elif isinstance(dt, datetime):
            naive = dt
        elif isinstance(dt, (int, float)):
            naive = datetime.fromtimestamp(dt)
        else:
            raise TypeError(f"Cannot construct SmartDateTime from {type(dt)}: {dt}")
        
        if naive.tzinfo is None:
            try:
                base = tz.localize(naive)   # pytz path
            except AttributeError:
                base = naive.replace(tzinfo=tz)  # zoneinfo path
        else: 
            base = naive.astimezone(tz)

        obj = super().__new__(
            cls,
            base.year, base.month, base.day,
            base.hour, base.minute, base.second,
            base.microsecond,
            tzinfo=base.tzinfo
        )

        obj.str_format = fmt
        obj.default_tz = tz
        return obj

    def __str__(self):
        """Format output according to str_format."""
        return self.strftime(self.str_format)
    
    # Arithmetic operations ---------------------------------------------------
    
    def __add__(self, b: Union[Dict, int, float, timedelta]):
        """
        Add interval (timedelta, dict params, or seconds) to datetime.
        Returns new SmartDateTime.
        """
        if isinstance(b, dict):
            delta = timedelta(**b)
        elif isinstance(b, (int, float)):
            delta = timedelta(seconds=b)
        elif isinstance(b, timedelta):
            delta = b
        else:
            raise TypeError(
                "Unsupported operand type(s) for +: 'SmartDateTime' and '{}'".format(type(b)))
        return super().__add__(delta)
    
    def __sub__(self, b):
        """
        Subtract interval or another date.
        If other is datetime or SmartDateTime, returns difference in Timedelta.
        If interval, returns new SmartDateTime.
        """
        if isinstance(b, dict):
            other = timedelta(**b)
        elif isinstance(b, (int, float)):
            other = timedelta(seconds=b)
        elif isinstance(b, timedelta):
            other = b
        elif isinstance(b, datetime):
            other = self._make_tz_aware(b)
        elif isinstance(b, SmartDateTime):
            other = b
        else:
            raise TypeError("Unsupported operand type(s) for -: 'SmartDateTime' and '{}'".format(type(b)))
        return super().__sub__(other)

    # Comparison operations -------------------------------------------------
    
    def __eq__(self, other: Any) -> bool:
        """
        Compare to numeric timestamp, datetime, or SmartDateTime.
        """
        if isinstance(other, int):
            return self.to_timestamp() == other
        if isinstance(other, float):
            return self.to_timestamp(as_int=False) == other
        
        if isinstance(other, SmartDateTime):
            return super().__eq__(other)
        if isinstance(other, datetime):
            return super().__eq__(self._make_tz_aware(other))

        return False
    
    def __ne__(self, other: Any) -> bool:
        """Negation of __eq__."""
        return not self == other

    def __lt__(self, other: Any) -> bool:
        """Less-than comparison."""
        if isinstance(other, SmartDateTime):
            return super().__lt__(other)

        if isinstance(other, datetime):
            # Caso other seja naïve, compare usando um datetime puro
            if other.tzinfo is None:
                naive_self = datetime(
                    self.year, self.month, self.day,
                    self.hour, self.minute, self.second,
                    self.microsecond
                )
                return naive_self < other

            # caso other seja aware, torne-o aware no mesmo TZ de self
            aware_other = self._make_tz_aware(other)
            return super().__lt__(aware_other)

        if isinstance(other, int):
            return self.to_timestamp() < other

        if isinstance(other, float):
            return self.to_timestamp(as_int=False) < other

        raise TypeError(
            f"Unsupported operand type(s) for '<': 'SmartDateTime' and '{type(other)}'"
        )

    def __le__(self, other: Any) -> bool:
        """Less-than or equal comparison."""
        if isinstance(other, SmartDateTime):
            return super().__le__(other)
        if isinstance(other, datetime):
            return super().__le__(self._make_tz_aware(other))
        if isinstance(other, int):
            return self.to_timestamp() <= other
        if isinstance(other, float):
            return self.to_timestamp(as_int=False) <= other
        raise TypeError(
                "Unsupported operand type(s) for '<=': 'SmartDateTime' and '{}'".format(type(other)))

    def __gt__(self, other: Any) -> bool:
        """Greater-than comparison."""
        if isinstance(other, SmartDateTime):
            return super().__gt__(other)
        if isinstance(other, datetime):
            return super().__gt__(self._make_tz_aware(other))
        if isinstance(other, int):
            return self.to_timestamp() > other
        if isinstance(other, float):
            return self.to_timestamp(as_int=False) > other
        raise TypeError(
                "Unsupported operand type(s) for '>': 'SmartDateTime' and '{}'".format(type(other)))

    def __ge__(self, other: Any) -> bool:
        """Greater-than or equal comparison."""
        if isinstance(other, SmartDateTime):
            return super().__ge__(other)
        if isinstance(other, datetime):
            return super().__ge__(self._make_tz_aware(other))
        if isinstance(other, int):
            return self.to_timestamp() >= other
        if isinstance(other, float):
            return self.to_timestamp(as_int=False) >= other
        raise TypeError(
                "Unsupported operand type(s) for '>=': 'SmartDateTime' and '{}'".format(type(other)))
    
    # Conversion methods -----------------------------------------------------
    
    def __int__(self) -> int:
        """Convert to integer timestamp."""
        return self.to_timestamp()

    def __float__(self) -> float:
        """Convert to float timestamp."""
        return self.to_timestamp(as_int=False)
    
    # Mapping protocol -------------------------------------------------------

    def __iter__(self):
        """Iterate over (key, value) pairs from to_dict()."""
        return iter(self.to_dict().items())
    
    def __len__(self):
        """Return number of keys in to_dict()."""
        return len(self.to_dict())

    # Remaining utilities inherited from SmartDateTimeMixin

# -------------------- Usage Examples --------------------
# Below are some usage examples for SmartDateTime:
#
# 1. Creating instances:
#   sd1 = SmartDateTime()  # now with default timezone (local time)
#   sd2 = SmartDateTime("2025-06-25 15:30:00") # with default timezone (local time)
#   sd3 = SmartDateTime(1625130000)  # from Unix timestamp
#   sd4 = SmartDateTime(year=2022, month=1, day=1, hour=12, minute=0)
#   sd5 = SmartDateTime(2022, 1, 1, 12, 0)
#   sd6 = SmartDateTime("2025.06.25 15:30", timezone="America/New_York", strftime="%Y.%m.%d %H:%M")
#
# 2. Basic arithmetic:
#   sd_next = sd2 + {"days": 3, "hours": 2}
#   sd_prev = sd2 - 3600  # subtract one hour
#
# 3. Timezone conversions:
#   sd_utc = sd2.to_utc()
#   sd_ny = sd2.to_timezone("America/New_York")
#
# 4. Business days:
#   sd_plus5 = sd2.add_business_days(5)
#   sd_minus3 = sd2.subtract_business_days(3)
#
# 5. Boundaries:
#   start_day = sd2.start_of_day()
#   end_month = sd2.end_of_month()
#
# 6. Information:
#   week = sd2.week_of_year()
#   quarter = sd2.quarter()
#   is_weekend = sd2.is_weekend()
#
# 7. Human-readable diff:
#   diff_str = sd2.humanize_diff(sd_next)  # e.g. 'in 3d 2h'
#
# 8. Integration with pandas:
#   import pandas as pd
#   df = pd.DataFrame({
#       'time': [SmartDateTime() + i*3600 for i in range(5)],
#       'value': range(5)
#   }, dtype='object') # or just omit the dtype='object' for pandas to infer datetime64
# 
#  # apply SmartDateTime operations:
#   df['dt'].apply(lambda x: x + {"hours": 1, "minutes": 55})
#
#   # convert column to datetime64:
#   df['time'] = pd.to_datetime(df['time'].apply(lambda x: x.to_iso()))