Convert random traffic to generator classes

pydtn/__init__.py

@@ -12,6 +12,7 @@ A simulation is made up from a Network which contains:
 __all__ = [
     'Network',
     'Buffer',
+    'PacketInfo',
 
     'Node',
     'EpidemicNode',
@@ -22,7 +23,11 @@ __all__ = [
     'random_trace',
 
     'Traffic',
-    'random_traffic',
+    'RandomManyToManyTraffic',
+    'RandomManyToOneTraffic',
+    'RandomOneToManyTraffic',
+    'RandomOneToOneTraffic',
+    'RandomTraffic',
 ]
 __version__ = '0.2'
 __author__ = 'Jarrod Pas <j.pas@usask.ca>'
@@ -174,36 +179,45 @@ class Network:
         return stats
 
 
+PacketInfo = namedtuple('PacketInfo', [
+    'source', 'destination', 'created', 'time_to_live', 'payload'
+])
+
+
 class Packet:
     """An item to route through the network."""
 
-    def __init__(self, network, traffic):
+    def __init__(self, network, info):
         """Create a packet within a network, based on traffic."""
         self.network = network
-        self._traffic = traffic
-
+        self.info = info
         self.recieved = None
         self._stats = defaultdict(int)
 
     @property
     def source(self):
-        """Return source node of the packet."""
-        return self.network.nodes[self._traffic.source]
+        """Source of packet."""
+        return self.network.nodes[self.info.source]
 
     @property
     def destination(self):
-        """Return destination node of the packet."""
-        return self.network.nodes[self._traffic.destination]
+        """Destination of packet."""
+        return self.network.nodes[self.info.destination]
 
     @property
     def created(self):
-        """Return created time of packet."""
-        return self._traffic.created
+        """Time of packet creation."""
+        return self.info.created
 
     @property
     def time_to_live(self):
-        """Return time to live of packet."""
-        return self._traffic.time_to_live
+        """Time to live of packet."""
+        return self.info.time_to_live
+
+    @property
+    def payload(self):
+        """Payload of packet."""
+        return self.info.payload
 
     @property
     def deadline(self):
@@ -240,7 +254,9 @@ class Packet:
 
     def __len__(self):
         """Return size of packet."""
-        return self._traffic.payload
+        if isinstance(self.payload, int):
+            return self.payload
+        return len(self.payload)
 
 
 class Buffer:
@@ -551,21 +567,101 @@ def csv_trace(path):
             yield Contact(now, source, destination, join)
 
 
-Traffic = namedtuple('Traffic', [
-    'source', 'destination', 'created', 'time_to_live', 'payload'
-])
+class Traffic:
+    """Base traffic generator."""
 
+    # pylint: disable=too-few-public-methods
 
-def random_traffic(nodes, start=0, speed=1, seed=None, **options):
-    """Generate traffic from random source to random destination every step."""
-    random = Random(seed)
+    def __init__(self, **options):
+        """Create traffic generator."""
+        self.time_to_live = options.get('time_to_live', float('inf'))
+        self.payload = options.get('payload', 0)
+
+    def __iter__(self):
+        """Yield no traffic."""
+        yield
+
+    def create_packet(self, now, source, destination):
+        """Create a packet info for the simulation."""
+        return PacketInfo(source=source,
+                          destination=destination,
+                          created=now,
+                          time_to_live=self.time_to_live,
+                          payload=self.payload)
+
+
+class RandomManyToManyTraffic(Traffic):
+    """Generate random traffic between nodes."""
+
+    def __init__(self, sources, destinations, **options):
+        """
+        Create a random traffic generator.
+
+        Option Arguments:
+        seed -- seed to use for the random number generator (default None)
+        start -- which time step the first packet is create (default 0)
+        step -- time between traffic generation (default 1)
+        traffic_per_step -- traffic to generate during each step (default 1)
+        time_to_live -- packet time to live (default infinite)
+        payload -- packet payload (default 0)
+        """
+        super().__init__(**options)
+
+        self.sources = list(sources)
+        self.destinations = list(destinations)
+
+        if len(self.sources) == 1 and self.sources == self.destinations:
+            raise ValueError('only source is the same as only destination')
+
+        self.seed = options.get('seed', None)
+        self.start = options.get('start', 0)
+        self.step = options.get('step', 1)
+        self.traffic_per_step = options.get('traffic_per_step', 1)
+
+    def random_pair(self, random):
+        """Return random pair of nodes."""
+        source = destination = random.choice(self.sources)
+        while destination == source:
+            destination = random.choice(self.destinations)
+        return source, destination
+
+    def __iter__(self):
+        """Yield infinite random traffic."""
+        random = Random(self.seed)
+
+        for now in count(start=self.start, step=self.step):
+            for _ in range(self.traffic_per_step):
+                source, destination = self.random_pair(random)
+                yield self.create_packet(now, source, destination)
+
+
+class RandomManyToOneTraffic(RandomManyToManyTraffic):
+    """Generate random traffic from any source to a single node."""
+
+    def __init__(self, sources, destination, **options):
+        """Create many to one traffic generator."""
+        super().__init__(sources, [destination], **options)
+
+
+class RandomOneToManyTraffic(RandomManyToManyTraffic):
+    """Generate random traffic from a single node to any destination."""
+
+    def __init__(self, source, destinations, **options):
+        """Create one to many traffic generator."""
+        super().__init__([source], destinations, **options)
+
+
+class RandomOneToOneTraffic(RandomManyToManyTraffic):
+    """Generate random traffic from a single node to a single node."""
+
+    def __init__(self, source, destination, **options):
+        """Create one to one traffic generator."""
+        super().__init__([source], [destination], **options)
 
-    if isinstance(nodes, dict):
-        nodes = list(nodes)
 
-    time_to_live = options.get('time_to_live', float('inf'))
-    payload = options.get('payload', 0)
+class RandomTraffic(RandomManyToManyTraffic):
+    """Generate random traffic between any node a group."""
 
-    for created in count(start=start, step=speed):
-        source, destination = random.sample(nodes, 2)
-        yield Traffic(source, destination, created, time_to_live, payload)
+    def __init__(self, nodes, **options):
+        """Create random traffic generator."""
+        super().__init__(nodes, nodes, **options)