Simu: Let Antennae copy a signal to set an argument

simulations/lib/location/antenna.py

@@ -1,6 +1,9 @@
 from functools import partial
+import copy
+from typing import Union
 
 from .location import Location
+from ..signals import Signal
 
 class Antenna(Location):
     """
@@ -17,18 +20,34 @@ class Antenna(Location):
     def __repr__(self):
         return "Antenna({}, {})".format(repr(self.x), repr(self.x))
 
-    def emit(self, signal: callable) -> callable:
+    def emit(self, signal: Union[Signal, callable]) -> Union[Signal, callable]:
         """
-        Return a function that emits a signal from the antenna's location
-        """
-        return partial(signal, x_0=self.x)
+        Emit signal from this antenna's location.
 
-    def recv(self, signal: callable) -> callable:
+        Note that this merely sets a default argument.
         """
-        Return a function that traces the signal as a function of time
-        at the antenna's location
+        if not isinstance(signal, Signal):
+            return partial(signal, x_0=self.x)
+        else:
+            new_signal = copy.copy(signal)
+            new_signal.x_0 = self.x
+
+            return new_signal
+
+    def recv(self, signal: Union[Signal, callable]) -> Union[Signal, callable]:
         """
-        return partial(signal, x_f=self.x)
+        Trace signal as a function of time at this antenna's
+        location.
+
+        Note that this merely sets a default argument.
+        """
+        if not isinstance(signal, Signal):
+            return partial(signal, x_f=self.x)
+        else:
+            new_signal = copy.copy(signal)
+            new_signal.x_f = self.x
+
+            return new_signal
 
     receive = recv