如何修改Python中datetime.__add__的行为并保留自定义子类实例？

嘿，我完全懂你碰到的这个坑——datetime是不可变类型，子类化之后那些内置方法（比如astimezone、__add__）默认会返回原生的datetime.datetime实例，根本不会用你的自定义子类。咱们来调整一下代码，让这些方法乖乖返回你的自定义类实例，同时保留你需要的DST偏移修正逻辑。

核心思路 #

要解决这个问题，关键在于重写所有会返回原生datetime的方法，在方法内部把原生实例转换成你的自定义子类实例。因为datetime是不可变的，所有创建新实例的操作都得通过__new__来，所以我们要确保每一步生成的新对象都是自定义datetime类的实例。

修改后的代码 #

#!/usr/bin/env python
import datetime as dt_orig
import pytz

class datetime(dt_orig.datetime):
    def __new__(cls, *args, **kwargs):
        # 处理传入时区字符串的情况（如果需要的话）
        tzstring = None
        if 'tzinfo' in kwargs and isinstance(kwargs['tzinfo'], str):
            tzstring = kwargs['tzinfo']
            kwargs['tzinfo'] = None
        
        # 先创建基础的自定义datetime实例
        dt_obj = super().__new__(cls, *args, **kwargs)
        
        # 如果传入了时区字符串，转换时区并确保返回的是自定义实例
        if tzstring is not None:
            tz = pytz.timezone(tzstring)
            localized = tz.localize(dt_obj)
            dt_obj = cls(
                localized.year, localized.month, localized.day,
                localized.hour, localized.minute, localized.second,
                localized.microsecond, tzinfo=localized.tzinfo
            )
        return dt_obj

    @classmethod
    def from_naive(cls, naive_dt, tzinfo):
        """辅助方法：把原生naive datetime转成自定义带时区的datetime"""
        if isinstance(tzinfo, str):
            tzinfo = pytz.timezone(tzinfo)
        localized = tzinfo.localize(naive_dt)
        return cls(
            localized.year, localized.month, localized.day,
            localized.hour, localized.minute, localized.second,
            localized.microsecond, tzinfo=localized.tzinfo
        )

    def __add__(self, other):
        if isinstance(other, dt_orig.timedelta):
            # 先调用原生__add__得到结果
            tmp = super().__add__(other)
            # 修正DST偏移问题
            corrected = tmp.astimezone(pytz.utc).astimezone(tmp.tzinfo)
            # 把修正后的原生datetime转成自定义实例
            return self.__class__(
                corrected.year, corrected.month, corrected.day,
                corrected.hour, corrected.minute, corrected.second,
                corrected.microsecond, tzinfo=corrected.tzinfo
            )
        else:
            return super().__add__(other)

    def __radd__(self, other):
        return self.__add__(other)

    def astimezone(self, tz=None):
        # 调用原生astimezone得到结果
        tmp = super().astimezone(tz)
        # 转成自定义实例
        return self.__class__(
            tmp.year, tmp.month, tmp.day,
            tmp.hour, tmp.minute, tmp.second,
            tmp.microsecond, tzinfo=tmp.tzinfo
        )

关键修改点说明 #

1. 重写astimezone方法：原生astimezone返回的是原生datetime，我们在方法最后把它转换成自定义子类实例，确保链式调用时始终返回自定义对象。
2. 修正__add__的返回值：原来的__add__里super().__add__返回原生对象，现在我们把修正DST偏移后的原生对象手动构造成自定义实例返回。
3. 时区字符串处理优化：原来的代码里astimezone会返回原生对象，现在改成手动构造自定义实例，避免丢失子类类型。
4. 添加辅助方法from_naive：方便把原生的naive datetime快速转换成自定义带时区的实例，减少重复代码。

测试一下 #

你可以用下面的代码验证效果：

# 创建带时区的自定义datetime
tz = pytz.timezone('Europe/Berlin')
dt = datetime(2023, 10, 29, 2, 30, tzinfo=tz)
print(f"原始对象类型: {type(dt)}")  # 应该是<class '__main__.datetime'>

# 加上timedelta
delta = dt_orig.timedelta(hours=1)
new_dt = dt + delta
print(f"相加后对象类型: {type(new_dt)}")  # 还是自定义类
print(f"原始时间: {dt}, 相加后时间: {new_dt}")
print(f"原始偏移量: {dt.utcoffset()}, 相加后偏移量: {new_dt.utcoffset()}")  # 偏移量应该正确处理DST

这样修改后，不管是调用astimezone还是做时间加减，返回的都是你的自定义datetime实例，同时保留了你需要的DST偏移修正逻辑。

内容的提问来源于stack exchange，提问作者Martin Thoma