Add pcap2csv CLI tool (#18)

* Add pcap2csv tool

* Add header row to CSV output

* Mock out test that fails on weekends

README.md

@@ -15,6 +15,28 @@ and does not endorse or recommend this library.
 
 ## Usage
 
+### pcap2csv
+
+You can use the included `pcap2csv` tool to create intraday minute bars from the pcap data files:
+
+```
+$ go install github.com/timpalpant/go-iex/pcap2csv
+$ pcap2csv < input.pcap > output.csv
+```
+
+which produces a CSV with OHLC data:
+
+```csv
+symbol,time,open,high,low,close,volume
+AAPL,2017-07-10T14:33:00Z,148.9100,149.0000,148.9100,148.9800,2527
+AMZN,2017-07-10T14:33:00Z,364.8900,364.9600,364.8900,364.9600,1486
+DISH,2017-07-10T14:33:00Z,49.6600,49.6600,49.6200,49.6200,1049
+ELLO,2017-07-10T14:33:00Z,10.0300,10.1100,10.0300,10.1100,3523
+FB,2017-07-10T14:33:00Z,46.4000,46.4000,46.3400,46.3400,1633
+FICO,2017-07-10T14:33:00Z,57.4200,57.4200,57.3500,57.3800,1717
+GOOD,2017-07-10T14:33:00Z,18.7700,18.7700,18.7300,18.7300,2459
+```
+
 ### pcap2json
 
 If you just need a tool to convert the provided pcap data files into JSON, you can use the included `pcap2json` tool:

client_test.go

@@ -114,7 +114,32 @@ func TestDEEP(t *testing.T) {
 }
 
 func TestBook(t *testing.T) {
-	c := setupTestClient()
+	body := `{
+		"YELP": {
+			"bids": [
+				{
+					"price": 63.09,
+					"size": 300,
+					"timestamp": 1494538496261
+				}
+			],
+			"asks": [
+				{
+					"price": 63.92,
+					"size": 300,
+					"timestamp": 1494538381896
+				},
+				{
+					"price": 63.97,
+					"size": 300,
+					"timestamp": 1494538381885
+				}
+			]
+		}
+	}`
+	httpc := mockHTTPClient{body: body, code: 200}
+	c := NewClient(&httpc)
+
 	symbols := []string{"SPY"}
 	result, err := c.GetBook(symbols)
 	if err != nil {

consolidator/bars.go

@@ -0,0 +1,81 @@
+package consolidator
+
+import (
+	"sort"
+	"time"
+
+	"github.com/timpalpant/go-iex/iextp/tops"
+)
+
+// Bar represents trades aggregated over a time interval.
+type Bar struct {
+	Symbol    string
+	OpenTime  time.Time
+	CloseTime time.Time
+	Open      float64
+	High      float64
+	Low       float64
+	Close     float64
+	Volume    int64
+}
+
+// Construct a Bar for each distinct symbol in the given list
+// of trades.
+func MakeBars(trades []*tops.TradeReportMessage) []*Bar {
+	bySymbol := groupTradesBySymbol(trades)
+	result := make([]*Bar, 0, len(bySymbol))
+	for _, trades := range bySymbol {
+		result = append(result, MakeBar(trades))
+	}
+
+	return result
+}
+
+// Construct a Bar from the given list of trades.
+// NOTE: Assumes all ticks are from the same symbol.
+func MakeBar(trades []*tops.TradeReportMessage) *Bar {
+	sort.Slice(trades, func(i, j int) bool {
+		return trades[i].Timestamp.Before(trades[j].Timestamp)
+	})
+
+	bar := &Bar{
+		Symbol:   trades[0].Symbol,
+		OpenTime: trades[0].Timestamp,
+	}
+
+	for _, trade := range trades {
+		updateBar(bar, trade)
+	}
+
+	return bar
+}
+
+func groupTradesBySymbol(trades []*tops.TradeReportMessage) map[string][]*tops.TradeReportMessage {
+	bySymbol := make(map[string][]*tops.TradeReportMessage)
+	for _, trade := range trades {
+		bySymbol[trade.Symbol] = append(bySymbol[trade.Symbol], trade)
+	}
+
+	return bySymbol
+}
+
+// Update the given bar to incorporate the trade.
+// Note this function assumes the security and times are compatible.
+func updateBar(bar *Bar, trade *tops.TradeReportMessage) {
+	price := trade.Price
+	if price > bar.High {
+		bar.High = price
+	}
+
+	if bar.Low == 0 || price < bar.Low {
+		bar.Low = price
+	}
+
+	if bar.Open == 0 {
+		bar.Open = price
+	}
+
+	bar.CloseTime = trade.Timestamp
+	bar.Close = price
+	bar.Volume += int64(trade.Size)
+}

pcap2csv/main.go

@@ -0,0 +1,116 @@
+// pcap2csv is a small binary for extracting IEXTP messages
+// from a pcap dump and converting them to minute-resolution bars
+// in CSV format for research.
+//
+// The pcap dump is read from stdin, and may be gzipped,
+// and the resulting CSV data is written to stdout.
+package main
+
+import (
+	"encoding/csv"
+	"io"
+	"log"
+	"os"
+	"sort"
+	"strconv"
+	"time"
+
+	"github.com/timpalpant/go-iex"
+	"github.com/timpalpant/go-iex/consolidator"
+	"github.com/timpalpant/go-iex/iextp/tops"
+)
+
+var header = []string{
+	"symbol",
+	"time",
+	"open",
+	"high",
+	"low",
+	"close",
+	"volume",
+}
+
+func makeBars(trades []*tops.TradeReportMessage, openTime, closeTime time.Time) []*consolidator.Bar {
+	bars := consolidator.MakeBars(trades)
+	for _, bar := range bars {
+		bar.OpenTime = openTime
+		bar.CloseTime = closeTime
+	}
+
+	sort.Slice(bars, func(i, j int) bool {
+		return bars[i].Symbol < bars[j].Symbol
+	})
+
+	return bars
+}
+
+func writeBar(bar *consolidator.Bar, w *csv.Writer) error {
+	row := []string{
+		bar.Symbol,
+		bar.OpenTime.Format(time.RFC3339),
+		strconv.FormatFloat(bar.Open, 'f', 4, 64),
+		strconv.FormatFloat(bar.High, 'f', 4, 64),
+		strconv.FormatFloat(bar.Low, 'f', 4, 64),
+		strconv.FormatFloat(bar.Close, 'f', 4, 64),
+		strconv.FormatInt(bar.Volume, 10),
+	}
+
+	return w.Write(row)
+}
+
+func writeBars(bars []*consolidator.Bar, w *csv.Writer) error {
+	for _, bar := range bars {
+		if err := writeBar(bar, w); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func main() {
+	packetSource, err := iex.NewPacketDataSource(os.Stdin)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	scanner := iex.NewPcapScanner(packetSource)
+	writer := csv.NewWriter(os.Stdout)
+	if err := writer.Write(header); err != nil {
+		log.Fatal(err)
+	}
+	defer writer.Flush()
+
+	var trades []*tops.TradeReportMessage
+	var openTime, closeTime time.Time
+	for {
+		msg, err := scanner.NextMessage()
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+
+			log.Fatal(err)
+		}
+
+		if msg, ok := msg.(*tops.TradeReportMessage); ok {
+			if openTime.IsZero() {
+				openTime = msg.Timestamp.Truncate(time.Minute)
+				closeTime = openTime.Add(time.Minute)
+			}
+
+			if msg.Timestamp.After(closeTime) && len(trades) > 0 {
+				bars := makeBars(trades, openTime, closeTime)
+				if err := writeBars(bars, writer); err != nil {
+					log.Fatal(err)
+				}
+
+				trades = trades[:0]
+				openTime = msg.Timestamp.Truncate(time.Minute)
+				closeTime = openTime.Add(time.Minute)
+			}
+
+			trades = append(trades, msg)
+		}
+	}
+}