Liskov'ify feed & merge.

zipline/components/feed.py

@@ -2,6 +2,7 @@ import logging
 from collections import Counter
 
 from zipline.core.component import Component
+from zipline.components.aggregator import Aggregate
 import zipline.protocol as zp
 
 from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE, \
@@ -9,7 +10,7 @@ from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE, \
 
 LOGGER = logging.getLogger('ZiplineLogger')
 
-class Feed(Component):
+class Feed(Aggregate):
     """
     Connects to N PULL sockets, publishing all messages received to a
     PUB socket. Published messages are guaranteed to be in chronological
@@ -35,71 +36,17 @@ class Feed(Component):
     def get_id(self):
         return "FEED"
 
-    @property
-    def get_type(self):
-        return COMPONENT_TYPE.CONDUIT
-
-    # -------------
-    # Core Methods
-    # -------------
+    # -------
+    # Sockets
+    # -------
 
     def open(self):
         self.pull_socket = self.bind_data()
         self.feed_socket = self.bind_feed()
 
-    def do_work(self):
-        # wait for synchronization reply from the host
-        socks = dict(self.poll.poll(self.heartbeat_timeout))
-
-        # TODO: Abstract this out, maybe on base component
-        if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
-            msg = self.control_in.recv()
-            event, payload = CONTROL_UNFRAME(msg)
-
-            # -- Heartbeat --
-            if event == CONTROL_PROTOCOL.HEARTBEAT:
-                # Heart outgoing
-                heartbeat_frame = CONTROL_FRAME(
-                    CONTROL_PROTOCOL.OK,
-                    payload
-                )
-                self.control_out.send(heartbeat_frame)
-
-            # -- Soft Kill --
-            elif event == CONTROL_PROTOCOL.SHUTDOWN:
-                self.signal_done()
-                self.shutdown()
-
-            # -- Hard Kill --
-            elif event == CONTROL_PROTOCOL.KILL:
-                self.kill()
-
-
-        if self.pull_socket in socks and socks[self.pull_socket] == self.zmq.POLLIN:
-            message = self.pull_socket.recv()
-
-            if message == str(CONTROL_PROTOCOL.DONE):
-                self.ds_finished_counter += 1
-
-                if len(self.data_buffer) == self.ds_finished_counter:
-                    #drain any remaining messages in the buffer
-                    LOGGER.debug("draining feed")
-                    self.drain()
-                    self.signal_done()
-            else:
-                try:
-                    event = self.unframe(message)
-                # deserialization error
-                except zp.INVALID_DATASOURCE_FRAME as exc:
-                    return self.signal_exception(exc)
-
-                try:
-                    self.append(event)
-                    self.send_next()
-
-                # Invalid message
-                except zp.INVALID_DATASOURCE_FRAME as exc:
-                    return self.signal_exception(exc)
+    # -------------
+    # Core Methods
+    # -------------
 
     def unframe(self, msg):
         return zp.DATASOURCE_UNFRAME(msg)
@@ -169,36 +116,3 @@ class Feed(Component):
 
         if earliest_event != None:
             return earliest_source.pop(0)
-
-    def is_full(self):
-        """
-        Indicates whether the buffer has messages in buffer for
-        all un-DONE, blocking sources.
-        """
-        for source_id, events in self.data_buffer.iteritems():
-            if len(events) == 0:
-                return False
-        return True
-
-    def pending_messages(self):
-        """
-        Returns the count of all events from all sources in the
-        buffer.
-        """
-        total = 0
-        for events in self.data_buffer.values():
-            total += len(events)
-        return total
-
-    def add_source(self, source_id):
-        """
-        Add a data source to the buffer.
-        """
-        self.data_buffer[source_id] = []
-
-    def __len__(self):
-        """
-        Buffer's length is same as internal map holding separate sorted
-        arrays of events keyed by source id.
-        """
-        return len(self.data_buffer)