bpo-32417: Make timedelta arithmetic respect subclasses (#10902)

* Make timedelta return subclass types

Previously timedelta would always return the `date` and `datetime`
types, regardless of what it is added to. This makes it return
an object of the type it was added to.

* Add tests for timedelta arithmetic on subclasses

* Make pure python timedelta return subclass types

* Add test for fromtimestamp with tz argument

* Add tests for subclass behavior in now

* Add news entry.

Fixes:
bpo-32417
bpo-35364

* More descriptive variable names in tests

Addresses Victor's comments

Lib/datetime.py

@@ -1014,7 +1014,7 @@ class date:
         if isinstance(other, timedelta):
             o = self.toordinal() + other.days
             if 0 < o <= _MAXORDINAL:
-                return date.fromordinal(o)
+                return type(self).fromordinal(o)
             raise OverflowError("result out of range")
         return NotImplemented
 
@@ -2024,7 +2024,7 @@ class datetime(date):
         hour, rem = divmod(delta.seconds, 3600)
         minute, second = divmod(rem, 60)
         if 0 < delta.days <= _MAXORDINAL:
-            return datetime.combine(date.fromordinal(delta.days),
+            return type(self).combine(date.fromordinal(delta.days),
                                       time(hour, minute, second,
                                            delta.microseconds,
                                            tzinfo=self._tzinfo))

Lib/test/datetimetester.py

@@ -820,6 +820,44 @@ class TestTimeDelta(HarmlessMixedComparison, unittest.TestCase):
         self.assertEqual(str(t3), str(t4))
         self.assertEqual(t4.as_hours(), -1)
 
+    def test_subclass_date(self):
+        class DateSubclass(date):
+            pass
+
+        d1 = DateSubclass(2018, 1, 5)
+        td = timedelta(days=1)
+
+        tests = [
+            ('add', lambda d, t: d + t, DateSubclass(2018, 1, 6)),
+            ('radd', lambda d, t: t + d, DateSubclass(2018, 1, 6)),
+            ('sub', lambda d, t: d - t, DateSubclass(2018, 1, 4)),
+        ]
+
+        for name, func, expected in tests:
+            with self.subTest(name):
+                act = func(d1, td)
+                self.assertEqual(act, expected)
+                self.assertIsInstance(act, DateSubclass)
+
+    def test_subclass_datetime(self):
+        class DateTimeSubclass(datetime):
+            pass
+
+        d1 = DateTimeSubclass(2018, 1, 5, 12, 30)
+        td = timedelta(days=1, minutes=30)
+
+        tests = [
+            ('add', lambda d, t: d + t, DateTimeSubclass(2018, 1, 6, 13)),
+            ('radd', lambda d, t: t + d, DateTimeSubclass(2018, 1, 6, 13)),
+            ('sub', lambda d, t: d - t, DateTimeSubclass(2018, 1, 4, 12)),
+        ]
+
+        for name, func, expected in tests:
+            with self.subTest(name):
+                act = func(d1, td)
+                self.assertEqual(act, expected)
+                self.assertIsInstance(act, DateTimeSubclass)
+
     def test_division(self):
         t = timedelta(hours=1, minutes=24, seconds=19)
         second = timedelta(seconds=1)
@@ -2604,33 +2642,58 @@ class TestDateTime(TestDate):
         ts = base_d.timestamp()
 
         test_cases = [
-            ('fromtimestamp', (ts,)),
+            ('fromtimestamp', (ts,), base_d),
             # See https://bugs.python.org/issue32417
-            # ('fromtimestamp', (ts, timezone.utc)),
+            ('fromtimestamp', (ts, timezone.utc),
-            ('utcfromtimestamp', (utc_ts,)),
+                               base_d.astimezone(timezone.utc)),
-            ('fromisoformat', (d_isoformat,)),
+            ('utcfromtimestamp', (utc_ts,), base_d),
-            ('strptime', (d_isoformat, '%Y-%m-%dT%H:%M:%S.%f')),
+            ('fromisoformat', (d_isoformat,), base_d),
-            ('combine', (date(*args[0:3]), time(*args[3:]))),
+            ('strptime', (d_isoformat, '%Y-%m-%dT%H:%M:%S.%f'), base_d),
+            ('combine', (date(*args[0:3]), time(*args[3:])), base_d),
         ]
 
-        for constr_name, constr_args in test_cases:
+        for constr_name, constr_args, expected in test_cases:
             for base_obj in (DateTimeSubclass, base_d):
                 # Test both the classmethod and method
                 with self.subTest(base_obj_type=type(base_obj),
                                   constr_name=constr_name):
-                    constr = getattr(base_obj, constr_name)
+                    constructor = getattr(base_obj, constr_name)
 
-                    dt = constr(*constr_args)
+                    dt = constructor(*constr_args)
 
                     # Test that it creates the right subclass
                     self.assertIsInstance(dt, DateTimeSubclass)
 
                     # Test that it's equal to the base object
-                    self.assertEqual(dt, base_d.replace(tzinfo=None))
+                    self.assertEqual(dt, expected)
 
                     # Test that it called the constructor
                     self.assertEqual(dt.extra, 7)
 
+    def test_subclass_now(self):
+        # Test that alternate constructors call the constructor
+        class DateTimeSubclass(self.theclass):
+            def __new__(cls, *args, **kwargs):
+                result = self.theclass.__new__(cls, *args, **kwargs)
+                result.extra = 7
+
+                return result
+
+        test_cases = [
+            ('now', 'now', {}),
+            ('utcnow', 'utcnow', {}),
+            ('now_utc', 'now', {'tz': timezone.utc}),
+            ('now_fixed', 'now', {'tz': timezone(timedelta(hours=-5), "EST")}),
+        ]
+
+        for name, meth_name, kwargs in test_cases:
+            with self.subTest(name):
+                constr = getattr(DateTimeSubclass, meth_name)
+                dt = constr(**kwargs)
+
+                self.assertIsInstance(dt, DateTimeSubclass)
+                self.assertEqual(dt.extra, 7)
+
     def test_fromisoformat_datetime(self):
         # Test that isoformat() is reversible
         base_dates = [

Misc/NEWS.d/next/Library/2018-12-04-13-35-36.bpo-32417._Y9SKM.rst

@@ -0,0 +1,6 @@
+Performing arithmetic between :class:`datetime.datetime` subclasses and
+:class:`datetime.timedelta` now returns an object of the same type as the
+:class:`datetime.datetime` subclass. As a result,
+:meth:`datetime.datetime.astimezone` and alternate constructors like
+:meth:`datetime.datetime.now` and :meth:`datetime.fromtimestamp` called with
+a ``tz`` argument now *also* retain their subclass.

Modules/_datetimemodule.c

@@ -3004,7 +3004,8 @@ add_date_timedelta(PyDateTime_Date *date, PyDateTime_Delta *delta, int negate)
     int day = GET_DAY(date) + (negate ? -deltadays : deltadays);
 
     if (normalize_date(&year, &month, &day) >= 0)
-        result = new_date(year, month, day);
+        result = new_date_subclass_ex(year, month, day,
+                                      (PyObject* )Py_TYPE(date));
     return result;
 }
 
@@ -5166,9 +5167,10 @@ add_datetime_timedelta(PyDateTime_DateTime *date, PyDateTime_Delta *delta,
         return NULL;
     }
 
-    return new_datetime(year, month, day,
+    return new_datetime_subclass_ex(year, month, day,
                                     hour, minute, second, microsecond,
-                        HASTZINFO(date) ? date->tzinfo : Py_None, 0);
+                                    HASTZINFO(date) ? date->tzinfo : Py_None,
+                                    (PyObject *)Py_TYPE(date));
 }
 
 static PyObject *