time.py

from typing import Optional, Tuple, Union, overload


_time_unit_map = {
    "ps": -12,
    "ns": -9,
    "us": -6,
    "ms": -3,
    "s": 0,
}

_time_repr_map = {v: k for k, v in _time_unit_map.items()}

_freq_unit_map = {
    "hz": 0,
    "khz": 3,
    "mhz": 6,
    "ghz": 9,
    "thz": 12,
}

_freq_repr_map = {
    0: "Hz",
    3: "kHz",
    6: "MHz",
    9: "GHz",
    12: "THz",
}


class Time:
    """"""

    def _normalize_to_smallest_unit(self, a: "Time", b: "Time") -> Tuple[int, int, int]:
        unit = min(a._unit, b._unit)
        a_scaled = a._scalar * (10 ** (a._unit - unit))
        b_scaled = b._scalar * (10 ** (b._unit - unit))
        return a_scaled, b_scaled, unit

    @classmethod
    def _make(cls, scalar: float, unit: int) -> "Time":
        assert unit in _time_repr_map
        self = cls.__new__(cls)
        self._init(scalar, unit)
        return self

    def _init(self, scalar: float, unit: int) -> None:
        self._scalar = scalar
        self._unit = unit

    def __init__(self, scalar: float, unit: str) -> None:
        self._init(scalar, _time_unit_map[unit.lower()])

    def in_units(self, unit: str) -> float:
        return self._scalar * (10 ** (self._unit - _time_unit_map[unit]))

    @property
    def fs(self) -> float:
        return self.in_units("fs")

    @property
    def ps(self) -> float:
        return self.in_units("ps")

    @property
    def ns(self) -> float:
        return self.in_units("ns")

    @property
    def us(self) -> float:
        return self.in_units("us")

    @property
    def ms(self) -> float:
        return self.in_units("ms")

    @property
    def s(self) -> float:
        return self.in_units("s")

    def __neg__(self) -> "Time":
        return Time._make(-self._scalar, self._unit)

    def __add__(self, other: "Time") -> "Time":
        self_norm_scalar, other_norm_scalar, unit = self._normalize_to_smallest_unit(
            self, other
        )
        return Time._make(self_norm_scalar + other_norm_scalar, unit)

    def __sub__(self, other: "Time") -> "Time":
        self_norm_scalar, other_norm_scalar, unit = self._normalize_to_smallest_unit(
            self, other
        )
        return Time._make(self_norm_scalar - other_norm_scalar, unit)

    def __mul__(self, other: float) -> "Time":
        if isinstance(other, (int, float)):
            return Time._make(self._scalar * other, self._unit)
        else:
            return NotImplemented

    def __rmul__(self, other: float) -> "Time":
        return self * other

    @overload
    def __truediv__(self, other: float) -> "Time":
        ...

    @overload
    def __truediv__(self, other: "Time") -> float:
        ...

    def __truediv__(self, other: Union[float, "Time"]) -> Union["Time", float]:
        if isinstance(other, (int, float)):
            return Time._make(self._scalar / other, self._unit)
        elif isinstance(other, Time):
            self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
                self, other
            )
            return self_norm_scalar / other_norm_scalar
        else:
            return NotImplemented

    def __rtruediv__(self, other: float) -> "Frequency":
        if isinstance(other, (int, float)):
            return Frequency._make(other / self._scalar, -self._unit)
        else:
            return NotImplemented

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self._scalar}, {_time_repr_map[self._unit]!r})"

    async def wait(self, round_mode: Optional[str] = None) -> None:
        from cocotb.triggers import Timer

        await Timer(self._scalar, self._unit, round_mode=round_mode)

    def __eq__(self, other: "Time") -> bool:
        if not isinstance(other, Time):
            return False
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar == other_norm_scalar

    def __hash__(self) -> int:
        return hash(self._scalar) + hash(self._unit)

    def __lt__(self, other: "Time") -> bool:
        if not isinstance(other, Time):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar < other_norm_scalar

    def __le__(self, other: "Time") -> bool:
        if not isinstance(other, Time):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar <= other_norm_scalar

    def __gt__(self, other: "Time") -> bool:
        if not isinstance(other, Time):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar > other_norm_scalar

    def __ge__(self, other: "Time") -> bool:
        if not isinstance(other, Time):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar >= other_norm_scalar


class Frequency:
    """"""

    def _normalize_to_smallest_unit(
        self, a: "Frequency", b: "Frequency"
    ) -> Tuple[int, int, int]:
        unit = min(a._unit, b._unit)
        a_scaled = a._scalar * (10 ** (a._unit - unit))
        b_scaled = b._scalar * (10 ** (b._unit - unit))
        return a_scaled, b_scaled, unit

    @classmethod
    def _make(cls, scalar: float, unit: int) -> "Time":
        assert unit in _freq_repr_map
        self = cls.__new__(cls)
        self._init(scalar, unit)
        return self

    def _init(self, scalar: float, unit: int) -> None:
        self._scalar = scalar
        self._unit = unit

    def __init__(self, scalar: float, unit: str) -> None:
        self._init(scalar, _freq_unit_map[unit.lower()])

    def in_units(self, unit: str) -> float:
        return self._scalar * (10 ** (self._unit - _freq_unit_map[unit]))

    @property
    def Hz(self) -> float:
        return self.in_units("hz")

    @property
    def kHz(self) -> float:
        return self.in_units("khz")

    @property
    def MHz(self) -> float:
        return self.in_units("mhz")

    @property
    def GHz(self) -> float:
        return self.in_units("ghz")

    @overload
    def __mul__(self, other: float) -> "Frequency":
        ...

    @overload
    def __mul__(self, other: Time) -> float:
        ...

    def __mul__(self, other: Union[float, Time]) -> Union["Frequency", float]:
        if isinstance(other, (int, float)):
            return Frequency._make(self._scalar * other, self._unit)
        elif isinstance(other, Time):
            scale = 10 ** (self._unit + other._unit)
            return self._scalar * other._scalar * scale
        else:
            return NotImplemented

    @overload
    def __rmul__(self, other: float) -> "Frequency":
        ...

    @overload
    def __rmul__(self, other: Time) -> float:
        ...

    def __rmul__(self, other: Union[float, Time]) -> Union["Frequency", float]:
        return self * other

    @overload
    def __truediv__(self, other: float) -> "Frequency":
        ...

    @overload
    def __truediv__(self, other: "Frequency") -> float:
        ...

    def __truediv__(
        self, other: Union[float, "Frequency"]
    ) -> Union["Frequency", float]:
        if isinstance(other, (int, float)):
            return Frequency._make(self._scalar / other, self._unit)
        elif isinstance(other, Frequency):
            self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
                self, other
            )
            return self_norm_scalar / other_norm_scalar
        else:
            return NotImplemented

    def __rtruediv__(self, other: float) -> Time:
        return Time._make(other / self._scalar, -self._unit)

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self._scalar}, {_freq_repr_map[self._unit]!r})"

    def __eq__(self, other: "Frequency") -> bool:
        if not isinstance(other, Frequency):
            return False
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar == other_norm_scalar

    def __hash__(self) -> int:
        return hash(self._scalar) + hash(self._unit)

    def __lt__(self, other: "Frequency") -> bool:
        if not isinstance(other, Frequency):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar < other_norm_scalar

    def __le__(self, other: "Frequency") -> bool:
        if not isinstance(other, Frequency):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar <= other_norm_scalar

    def __gt__(self, other: "Frequency") -> bool:
        if not isinstance(other, Frequency):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar > other_norm_scalar

    def __ge__(self, other: "Frequency") -> bool:
        if not isinstance(other, Frequency):
            return NotImplemented
        self_norm_scalar, other_norm_scalar, _ = self._normalize_to_smallest_unit(
            self, other
        )
        return self_norm_scalar >= other_norm_scalar


def test_time_units() -> None:
    assert Time(10, "ns").ps == 10000
    assert Time(1.678, "ms").ns == 1678000
    assert Time(123, "ns").us == 0.123
    assert Time(1, "s").ms == 1000
    assert Time(1, "ns").s == 1e-9


def test_freq_units() -> None:
    assert Frequency(1, "GHz").Hz == 1e9
    assert Frequency(0.123, "MHz").kHz == 123
    assert Frequency(123, "kHz").MHz == 0.123
    assert Frequency(67, "Hz").GHz == 67e-9


def test_time_compare() -> None:
    assert Time(10, "ns") == Time(10000, "ps")
    assert Time(10, "ns") < Time(11, "us")
    assert Time(10, "ns") <= Time(10, "ns")
    assert Time(1, "s") > Time(10, "ns")
    assert Time(11, "ms") >= Time(10, "ns")


def test_freq_compare() -> None:
    assert Frequency(1, "kHz") == Frequency(0.001, "MHz")
    assert Frequency(1, "Hz") < Frequency(1, "GHz")
    assert Frequency(1, "GHz") <= Frequency(1, "GHz")
    assert Frequency(1, "MHz") > Frequency(1, "kHz")
    assert Frequency(1, "kHz") >= Frequency(999, "Hz")


def test_time_scale() -> None:
    assert Time(10, "ns") * 2 == Time(20, "ns")
    assert 2 * Time(10, "ns") == Time(20, "ns")
    assert Time(10, "ns") / 2 == Time(5, "ns")


def test_freq_scale() -> None:
    assert Frequency(1, "MHz") * 1000 == Frequency(1, "GHz")
    assert 123 * Frequency(1, "kHz") == Frequency(123, "kHz")
    assert Frequency(1, "MHz") / 1000 == Frequency(1, "kHz")


def test_time_sum() -> None:
    assert Time(10, "ns") + Time(10, "ns") == Time(20, "ns")
    assert Time(500, "ns") - Time(1, "us") == Time(-500, "ns")
    assert -Time(100, "ps") == Time(-100, "ps")


def test_time_ratio() -> None:
    assert Time(10, "us") / Time(10, "ns") == 1000


def test_freq_ratio() -> None:
    assert Frequency(1, "GHz") / Frequency(123, "MHz") == 1000 / 123


def test_time_to_freq_and_back() -> None:
    assert (1 / Time(10, "ns")) == Frequency(1 / 10, "GHz")
    assert (1 / Frequency(1, "GHz")) == Time(1, "ns")


def test_time_mul_freq() -> None:
    assert (Time(10, "ns") * Frequency(1, "GHz")) == 10
    assert (Frequency(1.2, "GHz") * Time(17, "us")) == 1.2 * 17 * 1000