diff --git a/b_asic/signal_generator.py b/b_asic/signal_generator.py
index cf79873794d11435542d94ea0ac38ec198c4a4c4..7b25e586083d91366f24278c310227516e9e01e5 100644
--- a/b_asic/signal_generator.py
+++ b/b_asic/signal_generator.py
@@ -19,45 +19,53 @@ class SignalGenerator:
def __call__(self, time: int) -> complex:
raise NotImplementedError
- def __add__(self, other) -> "AddGenerator":
+ def __add__(self, other) -> "_AddGenerator":
if isinstance(other, Number):
- return AddGenerator(self, Constant(other))
- return AddGenerator(self, other)
+ return _AddGenerator(self, Constant(other))
+ if isinstance(other, SignalGenerator):
+ return _AddGenerator(self, other)
+ raise TypeError(f"Cannot add {other!r} to {type(self)}")
- def __radd__(self, other) -> "AddGenerator":
+ def __radd__(self, other) -> "_AddGenerator":
if isinstance(other, Number):
- return AddGenerator(self, Constant(other))
- return AddGenerator(self, other)
+ return _AddGenerator(Constant(other), self)
+ raise TypeError(f"Cannot add {type(self)} to {other!r}")
- def __sub__(self, other) -> "SubGenerator":
+ def __sub__(self, other) -> "_SubGenerator":
if isinstance(other, Number):
- return SubGenerator(self, Constant(other))
- return SubGenerator(self, other)
+ return _SubGenerator(self, Constant(other))
+ if isinstance(other, SignalGenerator):
+ return _SubGenerator(self, other)
+ raise TypeError(f"Cannot subtract {other!r} from {type(self)}")
- def __rsub__(self, other) -> "SubGenerator":
+ def __rsub__(self, other) -> "_SubGenerator":
if isinstance(other, Number):
- return SubGenerator(Constant(other), self)
- return SubGenerator(other, self)
+ return _SubGenerator(Constant(other), self)
+ raise TypeError(f"Cannot subtract {type(self)} from {other!r}")
- def __mul__(self, other) -> "MulGenerator":
+ def __mul__(self, other) -> "_MulGenerator":
if isinstance(other, Number):
- return MultGenerator(self, Constant(other))
- return MultGenerator(self, other)
+ return _MulGenerator(self, Constant(other))
+ if isinstance(other, SignalGenerator):
+ return _MulGenerator(self, other)
+ raise TypeError(f"Cannot multiply {type(self)} with {other!r}")
- def __rmul__(self, other) -> "MulGenerator":
+ def __rmul__(self, other) -> "_MulGenerator":
if isinstance(other, Number):
- return MultGenerator(self, Constant(other))
- return MultGenerator(self, other)
+ return _MulGenerator(Constant(other), self)
+ raise TypeError(f"Cannot multiply {other!r} with {type(self)}")
- def __truediv__(self, other) -> "MulGenerator":
+ def __truediv__(self, other) -> "_DivGenerator":
if isinstance(other, Number):
- return DivGenerator(self, Constant(other))
- return DivGenerator(self, other)
+ return _DivGenerator(self, Constant(other))
+ if isinstance(other, SignalGenerator):
+ return _DivGenerator(self, other)
+ raise TypeError(f"Cannot divide {type(self)} with {other!r}")
- def __rtruediv__(self, other) -> "MulGenerator":
+ def __rtruediv__(self, other) -> "_DivGenerator":
if isinstance(other, Number):
- return DivGenerator(Constant(other), self)
- return DivGenerator(other, self)
+ return _DivGenerator(Constant(other), self)
+ raise TypeError(f"Cannot divide {other!r} with {type(self)}")
class Impulse(SignalGenerator):
@@ -70,7 +78,7 @@ class Impulse(SignalGenerator):
The delay before the signal goes to 1 for one sample.
"""
- def __init__(self, delay: int = 0) -> Callable[[int], complex]:
+ def __init__(self, delay: int = 0) -> None:
self._delay = delay
def __call__(self, time: int) -> complex:
@@ -90,7 +98,7 @@ class Step(SignalGenerator):
The delay before the signal goes to 1.
"""
- def __init__(self, delay: int = 0) -> Callable[[int], complex]:
+ def __init__(self, delay: int = 0) -> None:
self._delay = delay
def __call__(self, time: int) -> complex:
@@ -110,7 +118,7 @@ class Constant(SignalGenerator):
The constant.
"""
- def __init__(self, constant: complex = 1.0) -> Callable[[int], complex]:
+ def __init__(self, constant: complex = 1.0) -> None:
self._constant = constant
def __call__(self, time: int) -> complex:
@@ -130,7 +138,7 @@ class ZeroPad(SignalGenerator):
The data that should be padded.
"""
- def __init__(self, data: Sequence[complex]) -> Callable[[int], complex]:
+ def __init__(self, data: Sequence[complex]) -> None:
self._data = data
self._len = len(data)
@@ -156,9 +164,7 @@ class Sinusoid(SignalGenerator):
The normalized phase offset.
"""
- def __init__(
- self, frequency: float, phase: float = 0.0
- ) -> Callable[[int], complex]:
+ def __init__(self, frequency: float, phase: float = 0.0) -> None:
self._frequency = frequency
self._phase = phase
@@ -173,95 +179,87 @@ class Sinusoid(SignalGenerator):
)
-class AddGenerator:
+class _AddGenerator(SignalGenerator):
"""
Signal generator that adds two signals.
"""
- def __init__(
- self, a: SignalGenerator, b: SignalGenerator
- ) -> Callable[[int], complex]:
+ def __init__(self, a: SignalGenerator, b: SignalGenerator) -> None:
self._a = a
self._b = b
def __call__(self, time: int) -> complex:
return self._a(time) + self._b(time)
- def __str__(self):
+ def __repr__(self):
return f"{self._a} + {self._b}"
-class SubGenerator:
+class _SubGenerator(SignalGenerator):
"""
Signal generator that subtracts two signals.
"""
- def __init__(
- self, a: SignalGenerator, b: SignalGenerator
- ) -> Callable[[int], complex]:
+ def __init__(self, a: SignalGenerator, b: SignalGenerator) -> None:
self._a = a
self._b = b
def __call__(self, time: int) -> complex:
return self._a(time) - self._b(time)
- def __str__(self):
+ def __repr__(self):
return f"{self._a} - {self._b}"
-class MultGenerator:
+class _MulGenerator(SignalGenerator):
"""
Signal generator that multiplies two signals.
"""
- def __init__(
- self, a: SignalGenerator, b: SignalGenerator
- ) -> Callable[[int], complex]:
+ def __init__(self, a: SignalGenerator, b: SignalGenerator) -> None:
self._a = a
self._b = b
def __call__(self, time: int) -> complex:
return self._a(time) * self._b(time)
- def __str__(self):
+ def __repr__(self):
a = (
f"({self._a})"
- if isinstance(self._a, (AddGenerator, SubGenerator))
+ if isinstance(self._a, (_AddGenerator, _SubGenerator))
else f"{self._a}"
)
b = (
f"({self._b})"
- if isinstance(self._b, (AddGenerator, SubGenerator))
+ if isinstance(self._b, (_AddGenerator, _SubGenerator))
else f"{self._b}"
)
return f"{a} * {b}"
-class DivGenerator:
+class _DivGenerator(SignalGenerator):
"""
Signal generator that divides two signals.
"""
- def __init__(
- self, a: SignalGenerator, b: SignalGenerator
- ) -> Callable[[int], complex]:
+ def __init__(self, a: SignalGenerator, b: SignalGenerator) -> None:
self._a = a
self._b = b
def __call__(self, time: int) -> complex:
return self._a(time) / self._b(time)
- def __str__(self):
+ def __repr__(self):
a = (
f"({self._a})"
- if isinstance(self._a, (AddGenerator, SubGenerator))
+ if isinstance(self._a, (_AddGenerator, _SubGenerator))
else f"{self._a}"
)
b = (
f"({self._b})"
if isinstance(
self._b,
- (AddGenerator, SubGenerator, MultGenerator, DivGenerator),
+ (_AddGenerator, _SubGenerator, _MulGenerator, _DivGenerator),
)
else f"{self._b}"
)
diff --git a/test/test_signal_generator.py b/test/test_signal_generator.py
index 085b555a597c0f5233d3d700b962f82080bfbda9..e11423bdca9dde748e7e9732d4b96035629a4655 100644
--- a/test/test_signal_generator.py
+++ b/test/test_signal_generator.py
@@ -2,7 +2,17 @@ from math import sqrt
import pytest
-from b_asic.signal_generator import Constant, Impulse, Sinusoid, Step, ZeroPad
+from b_asic.signal_generator import (
+ Constant,
+ Impulse,
+ Sinusoid,
+ Step,
+ ZeroPad,
+ _AddGenerator,
+ _DivGenerator,
+ _MulGenerator,
+ _SubGenerator,
+)
def test_impulse():
@@ -96,6 +106,7 @@ def test_addition():
assert g(3) == 0
assert str(g) == "Impulse() + Impulse(2)"
+ assert isinstance(g, _AddGenerator)
g = 1.0 + Impulse(2)
assert g(-1) == 1
@@ -104,7 +115,8 @@ def test_addition():
assert g(2) == 2
assert g(3) == 1
- assert str(g) == "Impulse(2) + 1.0"
+ assert str(g) == "1.0 + Impulse(2)"
+ assert isinstance(g, _AddGenerator)
g = Impulse(1) + 1.0
assert g(-1) == 1
@@ -114,6 +126,7 @@ def test_addition():
assert g(3) == 1
assert str(g) == "Impulse(1) + 1.0"
+ assert isinstance(g, _AddGenerator)
def test_subtraction():
@@ -125,6 +138,7 @@ def test_subtraction():
assert g(3) == 0
assert str(g) == "Impulse() - Impulse(2)"
+ assert isinstance(g, _SubGenerator)
g = 1.0 - Impulse(2)
assert g(-1) == 1
@@ -134,6 +148,7 @@ def test_subtraction():
assert g(3) == 1
assert str(g) == "1.0 - Impulse(2)"
+ assert isinstance(g, _SubGenerator)
g = Impulse(2) - 1.0
assert g(-1) == -1
@@ -143,6 +158,7 @@ def test_subtraction():
assert g(3) == -1
assert str(g) == "Impulse(2) - 1.0"
+ assert isinstance(g, _SubGenerator)
def test_multiplication():
@@ -153,6 +169,7 @@ def test_multiplication():
assert g(2) == 0
assert str(g) == "Impulse() * 0.5"
+ assert isinstance(g, _MulGenerator)
g = 2 * Sinusoid(0.5, 0.25)
assert g(0) == pytest.approx(sqrt(2))
@@ -160,7 +177,8 @@ def test_multiplication():
assert g(2) == pytest.approx(-sqrt(2))
assert g(3) == pytest.approx(-sqrt(2))
- assert str(g) == "Sinusoid(0.5, 0.25) * 2"
+ assert str(g) == "2 * Sinusoid(0.5, 0.25)"
+ assert isinstance(g, _MulGenerator)
g = Step(1) * (Sinusoid(0.5, 0.25) + 1.0)
assert g(0) == 0
@@ -169,6 +187,7 @@ def test_multiplication():
assert g(3) == pytest.approx(-sqrt(2) / 2 + 1)
assert str(g) == "Step(1) * (Sinusoid(0.5, 0.25) + 1.0)"
+ assert isinstance(g, _MulGenerator)
def test_division():
@@ -179,6 +198,7 @@ def test_division():
assert g(2) == 0.5
assert str(g) == "Step() / 2"
+ assert isinstance(g, _DivGenerator)
g = 0.5 / Step()
assert g(0) == 0.5
@@ -186,6 +206,7 @@ def test_division():
assert g(2) == 0.5
assert str(g) == "0.5 / Step()"
+ assert isinstance(g, _DivGenerator)
g = Sinusoid(0.5, 0.25) / (0.5 * Step())
assert g(0) == pytest.approx(sqrt(2))
@@ -193,4 +214,5 @@ def test_division():
assert g(2) == pytest.approx(-sqrt(2))
assert g(3) == pytest.approx(-sqrt(2))
- assert str(g) == "Sinusoid(0.5, 0.25) / (Step() * 0.5)"
+ assert str(g) == "Sinusoid(0.5, 0.25) / (0.5 * Step())"
+ assert isinstance(g, _DivGenerator)