track/backend/app/core/time_utils.py

"""全局北京时间 (UTC+8) 与工作日时长计算"""
from datetime import datetime, date, time, timedelta, timezone
from zoneinfo import ZoneInfo

BEIJING_TZ = ZoneInfo("Asia/Shanghai")


def get_beijing_time() -> datetime:
    """返回当前北京时间"""
    return datetime.now(BEIJING_TZ)


def to_beijing(dt: datetime | None) -> datetime | None:
    """将任意 datetime 统一转为北京时间 aware。

    - naive 时间按 UTC 处理（数据库 timestamptz 实存 UTC，SQLAlchemy 读出常为 naive）
    - 带时区时间直接 astimezone 到北京
    """
    if dt is None:
        return None
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc).astimezone(BEIJING_TZ)
    return dt.astimezone(BEIJING_TZ)


def working_duration_hours(
    start: datetime | None,
    end: datetime | None,
    holidays: set[date] | None = None,
) -> float:
    """计算 start~end 之间排除周末与节假日的工作小时数。

    - 周末（周六/周日）整天排除
    - holidays 中配置的放假日期整天排除
    - 其余日期按 24 小时连续计（一天内的时间都算）
    - start/end 可为 naive（按 UTC 转）或 aware 北京时间
    """
    if not holidays:
        holidays = set()
    s = to_beijing(start)
    e = to_beijing(end)
    if s is None or e is None or e <= s:
        return 0.0

    total = 0.0
    day = s.date()
    last = e.date()
    while day <= last:
        if day.weekday() < 5 and day not in holidays:
            seg_start = max(s, datetime.combine(day, time.min, tzinfo=BEIJING_TZ))
            seg_end = min(e, datetime.combine(day, time.max, tzinfo=BEIJING_TZ))
            if seg_end > seg_start:
                total += (seg_end - seg_start).total_seconds() / 3600
        day += timedelta(days=1)
    return round(total, 1)