Merge pull request #1 from pofallon/develop

Implement lookup with streams

Author: pofallon

commands/lookup.js

@@ -1,51 +1,43 @@
-const readline = require('readline')
+const through2 = require('through2')
+const ps = require('promise-streams')
+const split = require('split2')
+const parallel = require('parallel-transform')
+const from = require('from2-array')
 const JSONStream = require('JSONStream')
-const PQueue = require('p-queue')
 const CredentialManager = require('../lib/credential-manager')
 const Twitter = require('../lib/twitter')
-const BatchEmitter = require('../lib/batch-emitter')
+const batch = require('../lib/batch-stream')
 
-const lookup = {
-  async users (name, users, streams) {
-    await doLookup(name, '1.1/users/lookup.json?screen_name=', users, streams)
-  },
-  async statuses (name, ids, streams) {
-    await doLookup(name, '1.1/statuses/lookup.json?id=', ids, streams)
-  }
-}
-
-const doLookup = async function (name, api, items, streams = process) {
+const doLookup = async function (api, name, items, inout = process) {
   let creds = new CredentialManager(name)
   let [key, secret] = await creds.getKeyAndSecret('consumer')
   let twitter = new Twitter(key, secret)
   let [token, tokenSecret] = await creds.getKeyAndSecret('account')
   twitter.setToken(token, tokenSecret)
-  let queue = new PQueue({concurrency: 2})
-  let jsonStream = JSONStream.stringify()
-  jsonStream.pipe(streams.stdout)
-  await new Promise((resolve, reject) => {
-    let batch = new BatchEmitter(100)
-    batch.on('data', (data) => {
-      queue.add(() => twitter.get(`${api}${data.join(',')}`))
-        .then((results) => {
-          results.forEach((result) => { jsonStream.write(result) })
-        }).catch(reject)
-    })
-    batch.on('end', () => {
-      queue.onIdle().then(() => {
-        jsonStream.end()
-        resolve()
-      })
-    })
-    if (items) {
-      items.split(',').forEach((item) => { batch.add(item) })
-      batch.done()
-    } else {
-      readline.createInterface({input: streams.stdin})
-        .on('line', (line) => { batch.add(line) })
-        .on('close', () => { batch.done() })
-    }
-  })
+  return ps.pipeline(
+    items ? from.obj(items.split(',')) : inout.stdin.pipe(split()),
+    batch(100),
+    parallel(2, function (data, next) {
+      twitter.get(`${api}${data.join(',')}`)
+        .then((results) => next(null, results))
+        .catch(next)
+    }),
+    through2.obj(function (chunk, enc, next) {
+      chunk.forEach((c) => this.push(c))
+      next()
+    }),
+    JSONStream.stringify(),
+    inout.stdout
+  )
+}
+
+const lookup = {
+  async users (...args) {
+    await doLookup('1.1/users/lookup.json?screen_name=', ...args)
+  },
+  async statuses (...args) {
+    await doLookup('1.1/statuses/lookup.json?id=', ...args)
+  }
 }
 
 module.exports = lookup

lib/batch-emitter.js

@@ -0,0 +1,26 @@
+const through2 = require('through2')
+
+const batchStream = (batchSize = 10) => {
+  let batch = []
+  return through2.obj(
+    (chunk, enc, next) => {
+      batch.push(chunk)
+      if (batch.length === batchSize) {
+        let data = batch
+        batch = []
+        next(null, data)
+      } else {
+        next()
+      }
+    },
+    (next) => {
+      if (batch.length > 0) {
+        next(null, batch)
+      } else {
+        next()
+      }
+    }
+  )
+}
+
+module.exports = batchStream