Files
2026-07-13 13:05:14 +08:00

107 lines
3.6 KiB
Python

from __future__ import annotations
from datetime import datetime, timedelta, timezone
from typing import Any
import numpy as np
import pyarrow as pa
import pytest
from rerun.time import to_nanos, to_nanos_since_epoch
from typing_extensions import Self
VALID_TO_NANOS_CASES = [
(0, 0),
(2, 2_000_000_000),
(np.int8(3), 3_000_000_000),
(np.int16(4), 4_000_000_000),
(np.int32(5), 5_000_000_000),
(np.int64(6), 6_000_000_000),
(0.75, 750_000_000),
(np.float64(3.6), 3_600_000_000),
(timedelta(seconds=2, microseconds=500_000), 2_500_000_000),
(np.timedelta64(4, "s"), 4_000_000_000),
(np.timedelta64(2500, "ms"), 2_500_000_000),
(np.timedelta64(25, "ns"), 25),
(-3, -3_000_000_000),
(np.int64(-2), -2_000_000_000),
(np.float64(-1.5), -1_500_000_000),
(timedelta(seconds=-4, microseconds=-500_000), -4_500_000_000),
(np.timedelta64(-7, "s"), -7_000_000_000),
]
@pytest.mark.parametrize("duration,expected", VALID_TO_NANOS_CASES)
def test_to_nanos_valid(duration: int | float | timedelta | np.timedelta64, expected: int) -> None:
assert to_nanos(duration) == expected
INVALID_TO_NANOS_CASES = [
"invalid",
None,
[1, 2, 3],
3 + 4j,
]
@pytest.mark.parametrize("duration", INVALID_TO_NANOS_CASES)
def test_to_nanos_invalid(duration: Any) -> None:
with pytest.raises(TypeError):
to_nanos(duration)
class MockPandasTimestamp(datetime):
nanoseconds: int = 0
def __new__(cls, *args: Any, nanoseconds: int = 0, **kwargs: dict[str, Any]) -> Self:
instance = super().__new__(cls, *args, **kwargs) # type: ignore[arg-type]
object.__setattr__(instance, "nanoseconds", nanoseconds)
return instance
def to_datetime64(self) -> np.datetime64:
return np.datetime64(self.replace(microsecond=0).isoformat() + f".{self.nanoseconds:09d}", "ns")
VALID_TO_NANOS_SINCE_EPOCH_CASES = [
(0, 0),
(10, 10_000_000_000),
(np.int8(127), 127_000_000_000),
(np.int32(50), 50_000_000_000),
(np.int64(1223334444), 1_223_334_444_000_000_000),
(1.5, 1_500_000_000),
(np.float64(2.7), 2_700_000_000),
(datetime(1970, 1, 1), 0),
(datetime(1970, 1, 1, tzinfo=timezone.utc), 0),
(datetime(1970, 1, 1, 0, 0, 1), 1_000_000_000),
(np.datetime64("1970-01-01T00:00:00"), 0),
(np.datetime64("1970-01-01T00:00:01"), 1_000_000_000),
(np.datetime64("1970-01-01T00:00:00.000000000"), 0),
(datetime(1969, 12, 31, 23, 59, 59, tzinfo=timezone.utc), -1_000_000_000),
(np.datetime64("1969-12-31T23:59:59"), -1_000_000_000),
(datetime(2050, 1, 1, 0, 0, 0, tzinfo=timezone.utc), 2524608000000000000),
(np.datetime64("2050-01-01T00:00:00"), 2524608000000000000),
(np.datetime64("2000-01-01T00:00:00.123456789"), 946684800123456789),
(MockPandasTimestamp(2000, 1, 1, 0, 0, 0, 123456, nanoseconds=123456789), 946684800123456789),
(pa.scalar(np.datetime64("2000-01-01T00:00:00.123456789"), type=pa.timestamp("ns")), 946684800123456789),
]
@pytest.mark.parametrize("timestamp,expected", VALID_TO_NANOS_SINCE_EPOCH_CASES)
def test_to_nanos_since_epoch_valid(timestamp: int | float | datetime | np.datetime64, expected: int) -> None:
assert to_nanos_since_epoch(timestamp) == expected
INVALID_TO_NANOS_SINCE_EPOCH_CASES = [
"invalid",
None,
# Not enough precision for nanosecond scale
np.float16(14.234),
np.float32(211.623),
[1, 2, 3],
]
@pytest.mark.parametrize("timestamp", INVALID_TO_NANOS_SINCE_EPOCH_CASES)
def test_to_nanos_since_epoch_invalid(timestamp: Any) -> None:
with pytest.raises(TypeError):
to_nanos_since_epoch(timestamp)