fix: wip, added a bunch more tests

Author: ollz272

src/input/datetime.rs

@@ -109,6 +109,40 @@ impl<'a> EitherTimedelta<'a> {
             Self::Raw(duration) => duration_as_pytimedelta(py, duration),
         }
     }
+
+    pub fn total_seconds(&self) -> PyResult<f64> {
+        match self {
+            Self::Raw(timedelta) => {
+                let mut days: f64 = timedelta.day as f64;
+                let mut seconds: f64 = timedelta.second as f64;
+                let mut microseconds: f64 = timedelta.microsecond as f64;
+                let mut total_seconds: f64 = if !timedelta.positive {
+                    -1.0 * (86400.0 * days + seconds + microseconds / 1_000_000.0)
+                } else {
+                    86400.0 * days + seconds + microseconds / 1_000_000.0
+                };
+                Ok(total_seconds)
+            }
+            Self::PyExact(py_timedelta) => {
+                let mut days: f64 = py_timedelta.get_days() as f64; // -999999999 to 999999999
+                let mut seconds: f64 = py_timedelta.get_seconds() as f64; // 0 through 86399
+                let mut microseconds: f64 = py_timedelta.get_microseconds() as f64; // 0 through 999999
+                let positive = days >= 0.0;
+                let total_seconds: f64 = if !positive {
+                    86400.0 * days + seconds + microseconds / 1_000_000.0
+                } else {
+                    86400.0 * days + seconds + microseconds / 1_000_000.0
+                };
+                Ok(total_seconds)
+            }
+            Self::PySubclass(py_timedelta) => {
+                let total_seconds: f64 = py_timedelta
+                    .call_method0(intern!(py_timedelta.py(), "total_seconds"))?
+                    .extract()?;
+                Ok(total_seconds)
+            }
+        }
+    }
 }
 
 impl<'a> TryFrom<&'_ Bound<'a, PyAny>> for EitherTimedelta<'a> {

src/serializers/config.rs

@@ -110,10 +110,6 @@ serialization_mode! {
 }
 
 impl TimedeltaMode {
-    fn total_seconds<'py>(py_timedelta: &Bound<'py, PyDelta>) -> PyResult<Bound<'py, PyAny>> {
-        py_timedelta.call_method0(intern!(py_timedelta.py(), "total_seconds"))
-    }
-
     pub fn either_delta_to_json(self, py: Python, either_delta: &EitherTimedelta) -> PyResult<PyObject> {
         match self {
             Self::Iso8601 => {
@@ -123,15 +119,13 @@ impl TimedeltaMode {
             Self::Float | Self::SecondsFloat => {
                 // convert to int via a py timedelta not duration since we know this this case the input would have
                 // been a py timedelta
-                let py_timedelta = either_delta.try_into_py(py)?;
-                let seconds = Self::total_seconds(&py_timedelta)?;
+                let seconds: f64 = either_delta.total_seconds()?;
                 Ok(seconds.into_py(py))
             }
             Self::MillisecondsFloat => {
                 // convert to int via a py timedelta not duration since we know this this case the input would have
                 // been a py timedelta
-                let py_timedelta = either_delta.try_into_py(py)?;
-                let seconds: f64 = Self::total_seconds(&py_timedelta)?.extract()?;
+                let seconds: f64 = either_delta.total_seconds()?;
                 let object: Bound<PyFloat> = PyFloat::new_bound(py, seconds * 1000.0);
                 Ok(object.into_py(py))
             }
@@ -145,13 +139,11 @@ impl TimedeltaMode {
                 Ok(d.to_string().into())
             }
             Self::Float | Self::SecondsFloat => {
-                let py_timedelta = either_delta.try_into_py(py)?;
-                let seconds: f64 = Self::total_seconds(&py_timedelta)?.extract()?;
+                let seconds: f64 = either_delta.total_seconds()?;
                 Ok(seconds.to_string().into())
             }
             Self::MillisecondsFloat => {
-                let py_timedelta = either_delta.try_into_py(py)?;
-                let seconds: f64 = Self::total_seconds(&py_timedelta)?.extract()?;
+                let seconds: f64 = either_delta.total_seconds()?;
                 let milliseconds: f64 = seconds * 1000.0;
                 Ok(milliseconds.to_string().into())
             }
@@ -170,15 +162,11 @@ impl TimedeltaMode {
                 serializer.serialize_str(&d.to_string())
             }
             Self::Float | Self::SecondsFloat => {
-                let py_timedelta = either_delta.try_into_py(py).map_err(py_err_se_err)?;
-                let seconds = Self::total_seconds(&py_timedelta).map_err(py_err_se_err)?;
-                let seconds: f64 = seconds.extract().map_err(py_err_se_err)?;
+                let seconds: f64 = either_delta.total_seconds().map_err(py_err_se_err)?;
                 serializer.serialize_f64(seconds)
             }
             Self::MillisecondsFloat => {
-                let py_timedelta = either_delta.try_into_py(py).map_err(py_err_se_err)?;
-                let seconds = Self::total_seconds(&py_timedelta).map_err(py_err_se_err)?;
-                let seconds: f64 = seconds.extract().map_err(py_err_se_err)?;
+                let seconds: f64 = either_delta.total_seconds().map_err(py_err_se_err)?;
                 let milliseconds: f64 = seconds * 1000.0;
                 serializer.serialize_f64(milliseconds)
             }

tests/serializers/test_any.py

@@ -216,40 +216,54 @@ def test_any_config_timedelta_float_negative(mode):
     assert s.to_json({one_half_s: 'foo'}) == b'{"-1.5":"foo"}'
 
 
-def test_any_config_timedelta_millisecond():
-    s = SchemaSerializer(core_schema.any_schema(), config={'ser_json_timedelta': 'milliseconds_float'})
-    h2 = timedelta(hours=2)
-    assert s.to_python(h2) == h2
-    assert s.to_python(h2, mode='json') == 7200000.0
-    assert s.to_json(h2) == b'7200000.0'
-
-    assert s.to_python({h2: 'foo'}) == {h2: 'foo'}
-    assert s.to_python({h2: 'foo'}, mode='json') == {'7200000': 'foo'}
-    assert s.to_json({h2: 'foo'}) == b'{"7200000":"foo"}'
-
-
-def test_any_config_timedelta_millisecond_fraction():
-    s = SchemaSerializer(core_schema.any_schema(), config={'ser_json_timedelta': 'milliseconds_float'})
-    h2 = timedelta(seconds=1.5)
-    assert s.to_python(h2) == h2
-    assert s.to_python(h2, mode='json') == 1500.0
-    assert s.to_json(h2) == b'1500.0'
-
-    assert s.to_python({h2: 'foo'}) == {h2: 'foo'}
-    assert s.to_python({h2: 'foo'}, mode='json') == {'1500': 'foo'}
-    assert s.to_json({h2: 'foo'}) == b'{"1500":"foo"}'
-
-
-def test_any_config_timedelta_millisecond_negative():
+@pytest.mark.parametrize(
+    'td,expected_to_python,expected_to_json,expected_to_python_dict,expected_to_json_dict',
+    [
+        (timedelta(hours=2), 7200000.0, b'7200000.0', {'7200000': 'foo'}, b'{"7200000":"foo"}'),
+        (timedelta(hours=-2), -7200000.0, b'-7200000.0', {'-7200000': 'foo'}, b'{"-7200000":"foo"}'),
+        (timedelta(seconds=1.5), 1500.0, b'1500.0', {'1500': 'foo'}, b'{"1500":"foo"}'),
+        (timedelta(seconds=-1.5), -1500.0, b'-1500.0', {'-1500': 'foo'}, b'{"-1500":"foo"}'),
+        (timedelta(microseconds=1), 0.001, b'0.001', {'0.001': 'foo'}, b'{"0.001":"foo"}'),
+        (timedelta(microseconds=-1), -0.001, b'-0.001', {'-0.001': 'foo'}, b'{"-0.001":"foo"}'),
+        (timedelta(days=1), 86400000.0, b'86400000.0', {'86400000': 'foo'}, b'{"86400000":"foo"}'),
+        (timedelta(days=-1), -86400000.0, b'-86400000.0', {'-86400000': 'foo'}, b'{"-86400000":"foo"}'),
+        (timedelta(days=1, seconds=1), 86401000.0, b'86401000.0', {'86401000': 'foo'}, b'{"86401000":"foo"}'),
+        (timedelta(days=-1, seconds=-1), -86401000.0, b'-86401000.0', {'-86401000': 'foo'}, b'{"-86401000":"foo"}'),
+        (timedelta(days=1, seconds=-1), 86399000.0, b'86399000.0', {'86399000': 'foo'}, b'{"86399000":"foo"}'),
+        (
+            timedelta(days=1, seconds=1, microseconds=1),
+            86401000.001,
+            b'86401000.001',
+            {'86401000.001': 'foo'},
+            b'{"86401000.001":"foo"}',
+        ),
+        (
+            timedelta(days=-1, seconds=-1, microseconds=-1),
+            -86401000.001,
+            b'-86401000.001',
+            {'-86401000.001': 'foo'},
+            b'{"-86401000.001":"foo"}',
+        ),
+        (
+            timedelta(days=-1, seconds=-1, microseconds=-1),
+            -86401000.001,
+            b'-86401000.001',
+            {'-86401000.001': 'foo'},
+            b'{"-86401000.001":"foo"}',
+        ),
+    ],
+)
+def test_any_config_timedelta_millisecond(
+    td: timedelta, expected_to_python, expected_to_json, expected_to_python_dict, expected_to_json_dict
+):
     s = SchemaSerializer(core_schema.any_schema(), config={'ser_json_timedelta': 'milliseconds_float'})
-    h2 = timedelta(seconds=-1.5)
-    assert s.to_python(h2) == h2
-    assert s.to_python(h2, mode='json') == -1500.0
-    assert s.to_json(h2) == b'-1500.0'
+    assert s.to_python(td) == td
+    assert s.to_python(td, mode='json') == expected_to_python
+    assert s.to_json(td) == expected_to_json
 
-    assert s.to_python({h2: 'foo'}) == {h2: 'foo'}
-    assert s.to_python({h2: 'foo'}, mode='json') == {'-1500': 'foo'}
-    assert s.to_json({h2: 'foo'}) == b'{"-1500":"foo"}'
+    assert s.to_python({td: 'foo'}) == {td: 'foo'}
+    assert s.to_python({td: 'foo'}, mode='json') == expected_to_python_dict
+    assert s.to_json({td: 'foo'}) == expected_to_json_dict
 
 
 def test_recursion(any_serializer):