Repository: just-work/fffw

Issue:
TS: base class and magic methods
`TS` instance based on `timedelta` is not json serializable, what leads to complicated Celery setup.
1. TS may be based on float and implement datetime/timedelta support via  `__sub__`/`__add__` methods
2. TS should also define __repr__ method

Comments:

Patch:
diff --git a/fffw/graph/meta.py b/fffw/graph/meta.py
index ea0db03..93c0230 100644
--- a/fffw/graph/meta.py
+++ b/fffw/graph/meta.py
@@ -1,8 +1,13 @@
-import abc
 from dataclasses import dataclass
 from datetime import timedelta
 from enum import Enum
-from typing import List, Union, Any, Optional
+from functools import wraps
+from typing import List, Union, Any, Optional, Callable, overload, Tuple, cast
+
+try:
+    from typing import Literal
+except ImportError:  # pragma: no cover
+    from typing_extensions import Literal  # type: ignore
 
 from pymediainfo import MediaInfo  # type: ignore
 
@@ -29,8 +34,88 @@ class StreamType(Enum):
 VIDEO = StreamType.VIDEO
 AUDIO = StreamType.AUDIO
 
-
-class TS(timedelta):
+BinaryOp = Callable[[Any, Any], Any]
+BinaryTS = Callable[[Any, Any], "TS"]
+UnaryOp = Callable[[Any], Any]
+UnaryTS = Callable[[Any], "TS"]
+
+
+@overload
+def ts(func: BinaryOp,
+       *,
+       arg: bool = True,
+       res: Literal[True] = True,
+       noarg: Literal[False] = False
+       ) -> BinaryTS:
+    ...
+
+
+@overload
+def ts(func: BinaryOp,
+       *,
+       arg: bool = True,
+       res: Literal[False],
+       noarg: Literal[False] = False
+       ) -> BinaryOp:
+    ...
+
+
+@overload
+def ts(func: UnaryOp,
+       *,
+       arg: bool = True,
+       res: Literal[True] = True,
+       noarg: Literal[True]
+       ) -> UnaryTS:
+    ...
+
+
+@overload
+def ts(func: UnaryOp,
+       *,
+       arg: bool = True,
+       res: Literal[False],
+       noarg: Literal[True]
+       ) -> UnaryOp:
+    ...
+
+
+def ts(func: Union[BinaryOp, UnaryOp], *,
+       arg: bool = True, res: bool = True, noarg: bool = False
+       ) -> Union[BinaryOp, UnaryOp]:
+    """
+    Decorates functions to automatically cast first argument and result to TS.
+    """
+    if arg and res:
+        if noarg:
+            @wraps(func)
+            def wrapper(self: "TS") -> "TS":
+                return TS(cast(UnaryOp, func)(self))  # noqa
+        else:
+            @wraps(func)
+            def wrapper(self: "TS", value: Any) -> "TS":
+                if value is None:
+                    res = cast(BinaryOp, func)(self, value)  # noqa
+                else:
+                    res = cast(BinaryOp, func)(self, TS(value))  # noqa
+                return TS(res)
+    elif arg:
+        @wraps(func)
+        def wrapper(self: "TS", value: Any) -> Any:
+            if value is None:
+                return cast(BinaryOp, func)(self, value)  # noqa
+            return cast(BinaryOp, func)(self, TS(value))  # noqa
+    elif res:
+        @wraps(func)
+        def wrapper(self: "TS", value: Any) -> "TS":
+            return TS(cast(BinaryOp, func)(self, value))  # noqa
+    else:
+        return func
+
+    return wrapper
+
+
+class TS(float):
     """
     Timestamp data type.
 
@@ -38,18 +123,15 @@ class TS(timedelta):
     Integer values are parsed as milliseconds.
     """
 
-    def __new__(cls, value: Union[int, float, str], *args: int) -> "TS":
+    def __new__(cls, value: Union[int, float, str, timedelta]) -> "TS":
         """
         :param value: integer duration in milliseconds, float duration in
             seconds or string ffmpeg interval definition (123:59:59.999).
         :returns new timestamp from value.
         """
-        if args:
-            # from deconstruction
-            if not isinstance(value, int):
-                raise ValueError(value)
-            value = timedelta(value, *args).total_seconds()
-        if isinstance(value, int):
+        if isinstance(value, timedelta):
+            value = value.total_seconds()
+        elif isinstance(value, int):
             value = value / 1000.0
         elif isinstance(value, str):
             if '.' in value:
@@ -62,13 +144,86 @@ class TS(timedelta):
                 seconds *= 60
                 seconds += part
             value = seconds + fractional
-        return super().__new__(cls, seconds=value)  # type: ignore
+        return super().__new__(cls, value)  # type: ignore
+
+    __add__ = ts(float.__add__)
+    __radd__ = ts(float.__radd__)
+    __sub__ = ts(float.__sub__)
+    __rsub__ = ts(float.__rsub__)
+    __mul__ = ts(float.__mul__, arg=False)
+    __rmul__ = ts(float.__rmul__, arg=False)
+    __neg__ = ts(float.__neg__, noarg=True)
+    __abs__ = ts(float.__abs__, noarg=True)
+    __eq__ = ts(float.__eq__, res=False)
+    __ne__ = ts(float.__ne__, res=False)
+    __gt__ = ts(float.__gt__, res=False)
+    __ge__ = ts(float.__ge__, res=False)
+    __lt__ = ts(float.__lt__, res=False)
+    __le__ = ts(float.__le__, res=False)
+
+    @overload  # type: ignore
+    def __floordiv__(self, other: "TS") -> int:
+        ...
+
+    @overload  # type: ignore
+    def __floordiv__(self, other: int) -> "TS":
+        ...
+
+    def __floordiv__(self, other: Union["TS", float, int]) -> Union[int, "TS"]:
+        """
+        Division behavior from timedelta (rounds to microseconds)
+
+        >>> TS(10.0) // TS(3.0)
+        3
+        >>> TS(10.0) // 3
+        TS(3.333333)
+        >>> TS(10.0) // 3.0
+        TS(3.333333)
+        """
+        value = (float(self * 1000000.0) // other) / 1000000.0
+        if isinstance(other, TS):
+            return int(value)
+        return TS(value)
+
+    @overload
+    def __truediv__(self, other: "TS") -> float:  # type: ignore
+        ...
+
+    @overload
+    def __truediv__(self, other: Union[float, int]) -> "TS":  # type: ignore
+        ...
 
-    def __float__(self) -> float:
+    def __truediv__(self, other: Union["TS", float, int]) -> Union[float, "TS"]:
         """
-        :returns: duration in seconds.
+        Division behavior from timedelta
+
+        >>> TS(10.0) / TS(2.125)
+        4.705882352941177
+        >>> TS(10.0) / 2.125
+        TS(4.705882352941177)
+        >>> TS(10.0) / 2
+        TS(5.0)
+        """
+        value = super().__truediv__(other)
+        if isinstance(other, TS):
+            return value
+        return TS(value)
+
+    def __divmod__(self, other: float) -> Tuple[int, "TS"]:
         """
-        return self.total_seconds()
+        Div/mod behavior from timedelta
+
+        >>> divmod(TS(10.0), TS(2.125))
+        (4, TS(1.5))
+        """
+        div, mod = super().__divmod__(other)
+        return int(div), TS(mod)
+
+    def __int__(self) -> int:
+        """
+        :return: duration in milliseconds.
+        """
+        return int(float(self * 1000))
 
     def __str__(self) -> str:
         """
@@ -76,25 +231,30 @@ class TS(timedelta):
 
         :returns: ffmpeg seconds definition (123456.999).
         """
-        v = str(self.total_seconds())
+        v = super().__repr__()
         if '.' in v:
             v = v.rstrip('0')
+            if v.endswith('.'):
+                v += '0'
         return v
 
-    def __add__(self, other: timedelta) -> "TS":
-        if not isinstance(other, timedelta):
-            return NotImplemented
-        return TS(self.total_seconds() + other.total_seconds())
+    def __repr__(self) -> str:
+        return f'TS({super().__repr__()})'
 
-    def __sub__(self, other: timedelta) -> "TS":
-        if not isinstance(other, timedelta):
-            return NotImplemented
-        return TS(self.total_seconds() - other.total_seconds())
+    def total_seconds(self) -> float:
+        return float(self)
 
-    def __lt__(self, other: Union[int, float, str, timedelta, "TS"]) -> bool:
-        if not isinstance(other, timedelta):
-            other = TS(other)
-        return self.total_seconds() < other.total_seconds()
+    @property
+    def days(self) -> int:
+        return int(float(self / (24 * 3600)))
+
+    @property
+    def seconds(self) -> int:
+        return int(float(self) % (24 * 3600))
+
+    @property
+    def microseconds(self) -> int:
+        return int(float(self * 1000000) % 1000000)
 
 
 @dataclass


Test patch:
diff --git a/tests/test_meta_data.py b/tests/test_meta_data.py
index 0e00b0d..c6087db 100644
--- a/tests/test_meta_data.py
+++ b/tests/test_meta_data.py
@@ -1,6 +1,9 @@
+import json
 from copy import deepcopy
+from datetime import timedelta
+from typing import Iterable, Tuple, Any
 from unittest import TestCase
-
+from itertools import product
 from pymediainfo import MediaInfo  # type: ignore
 
 from fffw.graph import meta
@@ -206,23 +209,6 @@ class MetaDataTestCase(TestCase):
     def setUp(self) -> None:
         self.media_info = MediaInfo(SAMPLE)
 
-    def test_ts_deconstruction(self):
-        ts = meta.TS(3600 * 24 * 2 + 3600 * 4 + 0.123)
-        self.assertEqual(ts, deepcopy(ts))
-
-    def test_ts_init(self):
-        self.assertEqual(float(meta.TS(1.23)), 1.23)
-        self.assertEqual(float(meta.TS(1)), 0.001)
-        self.assertEqual(float(meta.TS('01:02:03.04')),
-                         1 * 3600 + 2 * 60 + 3 + 0.04)
-        self.assertEqual(float(meta.TS(1, 2, 3)),
-                         1 * 24 * 3600 + 2 + 0.000003)
-        self.assertRaises(ValueError, meta.TS, 1.1, 2, 3)
-
-    def test_ts_float(self):
-        ts = meta.TS(3600 * 24 * 2 + 3600 * 4 + 0.123)
-        self.assertEqual(float(ts), ts.total_seconds())
-
     def test_parse_streams(self):
         streams = meta.from_media_info(self.media_info)
         self.assertEqual(len(streams), 2)
@@ -260,3 +246,169 @@ class MetaDataTestCase(TestCase):
             samples=323616,
         )
         self.assertEqual(expected, audio)
+
+
+class TimeStampTestCase(TestCase):
+    td: timedelta
+    ts: meta.TS
+    binary_cases: Iterable[Tuple[Any, Any]]
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.td = timedelta(
+            days=10,
+            hours=12,
+            minutes=34,
+            seconds=56,
+            microseconds=789000)
+        cls.ts = meta.TS(cls.td.total_seconds())
+
+        ms = int(cls.td.total_seconds() * 1000)
+        seconds = cls.td.total_seconds()
+        string = '252:34:56.789000'
+        cls.binary_cases = (
+            (cls.ts, cls.ts),
+            (cls.ts, cls.td),
+            (cls.ts, ms),
+            (cls.ts, seconds),
+            (cls.ts, string),
+            (cls.td, cls.ts),
+            (ms, cls.ts),
+            (seconds, cls.ts),
+            (string, cls.ts),
+        )
+
+    def assert_ts_equal(self, ts: meta.TS, expected: float):
+        self.assertIsInstance(ts, meta.TS)
+        self.assertAlmostEqual(ts.total_seconds(), expected, places=4)
+
+    def test_ts_float(self):
+        self.assertEqual(float(self.ts), self.td.total_seconds())
+
+    def test_ts_int(self):
+        self.assertEqual(int(self.ts), int(self.td.total_seconds() * 1000))
+
+    def test_ts_deconstruction(self):
+        self.assertEqual(self.ts, deepcopy(self.ts))
+
+    def test_ts_init(self):
+        cases = (
+            # from float seconds
+            self.td.total_seconds(),
+            # from in milliseconds
+            int(self.td.total_seconds() * 1000),
+            # from string
+            '252:34:56.789000',
+        )
+        for v in cases:
+            with self.subTest(v):
+                self.assertEqual(self.ts, meta.TS(v))
+
+    def test_addition(self):
+        for case in self.binary_cases:
+            with self.subTest(case):
+                first, second = case
+                ts = first + second
+                self.assert_ts_equal(ts, 2 * self.td.total_seconds())
+
+    def test_substraction(self):
+        for case in self.binary_cases:
+            with self.subTest(case):
+                first, second = case
+                ts = first - second
+                self.assert_ts_equal(ts, 0.0)
+
+    def test_multiplication(self):
+        cases = (
+            2.0,
+            2,
+        )
+        for case in product(cases, (True, False)):
+            with self.subTest(case):
+                v, rev = case
+                if rev:
+                    ts = v * self.ts
+                else:
+                    ts = self.ts * v
+                self.assert_ts_equal(ts, 2 * self.td.total_seconds())
+
+    def test_divmod(self):
+        """ Test timedelta.__divmod__ behavior."""
+        # noinspection PyTypeChecker
+        div, mod = divmod(self.ts, self.ts)
+        self.assert_ts_equal(mod, 0.0)
+        self.assertIsInstance(div, int)
+        self.assertEqual(div, 1)
+
+    def test_floordiv(self):
+        """ Test timedelta.__floordiv__ behavior."""
+        ts = (self.ts + 0.000001) // 2
+        expected = int(self.td.total_seconds() * 1000000) / 2000000.0
+        self.assert_ts_equal(ts, expected)
+
+        ts = (self.ts + 0.000001) // meta.TS(2.0)
+        expected = int(self.td.total_seconds() * 1000000) // 2000000
+        self.assertIsInstance(ts, int)
+        self.assertEqual(ts, expected)
+
+    def test_truediv(self):
+        ts = (self.ts + 0.000001) / 2
+        expected = int(self.td.total_seconds() * 1000000) / 2000000.0
+        self.assert_ts_equal(ts, expected)
+
+        ts = (self.ts + 0.000001) / 2.0
+        expected = int(self.td.total_seconds() * 1000000) / 2000000.0
+        self.assert_ts_equal(ts, expected)
+
+        ts = (self.ts + 0.000001) / meta.TS(2.0)
+        expected = int(self.td.total_seconds() * 1000000) / 2000000.0
+        self.assertIsInstance(ts, float)
+        self.assertAlmostEqual(ts, expected, places=5)
+
+    def test_negate_abs(self):
+        ts = -self.ts
+        self.assert_ts_equal(ts, -self.td.total_seconds())
+        self.assert_ts_equal(abs(ts), self.td.total_seconds())
+
+    def test_compare(self):
+        v = self.ts + 0.001
+        cases = (
+            v,
+            v.total_seconds(),
+            int(v.total_seconds() * 1000),
+        )
+        for v in cases:
+            with self.subTest(v):
+                self.assertTrue(v > self.ts)
+                self.assertTrue(v >= self.ts)
+                self.assertFalse(v < self.ts)
+                self.assertTrue(self.ts < v)
+                self.assertTrue(self.ts <= v)
+                self.assertFalse(self.ts > v)
+                self.assertFalse(self.ts == v)
+                self.assertFalse(v == self.ts)
+                self.assertTrue(v != self.ts)
+                self.assertTrue(self.ts != v)
+
+        self.assertFalse(self.ts == None)  # noqa
+        self.assertTrue(self.ts != None)  # noqa
+        self.assertFalse(self.ts is None)
+        self.assertTrue(self.ts is not None)
+
+    def test_total_seconds(self):
+        self.assertEqual(self.ts.total_seconds(), self.td.total_seconds())
+
+    def test_fields(self):
+        self.assertEqual(self.ts.days, self.td.days)
+        self.assertEqual(self.ts.seconds, self.td.seconds)
+        self.assertEqual(self.ts.microseconds, self.td.microseconds)
+
+    def test_json_serializable(self):
+        self.assertEqual(json.dumps(self.ts),
+                         json.dumps(self.td.total_seconds()))
+
+    def test_str(self):
+        self.assertEqual(str(self.ts), str(self.td.total_seconds()))
+
+    def test_repr(self):
+        self.assertEqual(repr(self.ts), f'TS({repr(self.td.total_seconds())})')