Moved screenshot functions into Device's class.

Summary:
I have moved the screenshot functions from ScreenCaptureButtons to the Device classes. I have slightly rewritten them so that they return a Promise which resolves to a Buffer. The Buffer can then be saved to a file or converted to a data Blob.

I have removed streaming and simply loaded the image into memory. Once the image is in memory it can be manipulated for various tasks i.e. written to a file, or displayed in the app.

iOS screenshots had to be rewritten. I now save the image to a temp folder, load it into the apps memory, and then remove the temp image.

Reviewed By: jknoxville

Differential Revision: D16939901

fbshipit-source-id: 3e39a5aeda8d48829ac5a8ff912a98f110341c07

src/chrome/ScreenCaptureButtons.tsx

@@ -5,7 +5,7 @@
  * @format
  */
 
-import {Button, ButtonGroup} from 'flipper';
+import {Button, ButtonGroup, writeBufferToFile} from 'flipper';
 import React, {Component} from 'react';
 import {connect} from 'react-redux';
 import AndroidDevice from '../devices/AndroidDevice';
@@ -26,17 +26,10 @@ const CAPTURE_LOCATION = expandTilde(
   config().screenCapturePath || remote.app.getPath('desktop'),
 );
 
-type PullTransfer = any;
-
 type OwnProps = {};
 
 type StateFromProps = {
-  selectedDevice:
-    | BaseDevice
-    | typeof AndroidDevice
-    | typeof IOSDevice
-    | null
-    | undefined;
+  selectedDevice: BaseDevice | null | undefined;
 };
 
 type DispatchFromProps = {};
@@ -73,17 +66,6 @@ function getFileName(extension: 'png' | 'mp4'): string {
   return `Screen Capture ${new Date().toISOString()}.${extension}`;
 }
 
-function writePngStreamToFile(stream: PullTransfer): Promise<string> {
-  return new Promise((resolve, reject) => {
-    const pngPath = path.join(CAPTURE_LOCATION, getFileName('png'));
-    stream.on('end', () => {
-      resolve(pngPath);
-    });
-    stream.on('error', reject);
-    stream.pipe(fs.createWriteStream(pngPath));
-  });
-}
-
 type Props = OwnProps & StateFromProps & DispatchFromProps;
 class ScreenCaptureButtons extends Component<Props, State> {
   iOSRecorder: any | null | undefined;
@@ -132,32 +114,14 @@ class ScreenCaptureButtons extends Component<Props, State> {
 
   captureScreenshot: Promise<void> | any = () => {
     const {selectedDevice} = this.props;
-
-    if (selectedDevice instanceof AndroidDevice) {
-      return reportPlatformFailures(
-        selectedDevice.adb
-          .screencap(selectedDevice.serial)
-          .then(writePngStreamToFile)
-          .then(openFile),
-        'captureScreenshotAndroid',
-      ).catch(console.error);
-    } else if (selectedDevice instanceof IOSDevice) {
-      const screenshotPath = path.join(CAPTURE_LOCATION, getFileName('png'));
-      return reportPlatformFailures(
-        new Promise((resolve, reject) => {
-          exec(
-            `xcrun simctl io booted screenshot "${screenshotPath}"`,
-            async err => {
-              if (err) {
-                reject(err);
-              } else {
-                openFile(screenshotPath);
-                resolve();
-              }
-            },
-          );
-        }),
-        'captureScreenshotIos',
+    const pngPath = path.join(CAPTURE_LOCATION, getFileName('png'));
+    if (selectedDevice != null) {
+      reportPlatformFailures(
+        selectedDevice
+          .screenshot()
+          .then((buffer: Buffer) => writeBufferToFile(pngPath, buffer))
+          .then((path: string) => openFile(path)),
+        'captureScreenshot',
       );
     }
   };

src/devices/AndroidDevice.tsx

@@ -10,6 +10,7 @@ import {Priority} from 'adbkit-logcat-fb';
 import child_process from 'child_process';
 import child_process_promise from 'child-process-es6-promise';
 import ArchivedDevice from './ArchivedDevice';
+import {ReadStream} from 'fs';
 
 type ADBClient = any;
 
@@ -99,4 +100,18 @@ export default class AndroidDevice extends BaseDevice {
     const shellCommand = `am start ${encodeURI(location)}`;
     this.adb.shell(this.serial, shellCommand);
   }
+
+  screenshot(): Promise<Buffer> {
+    return new Promise((resolve, reject) => {
+      this.adb.screencap(this.serial).then((stream: ReadStream) => {
+        const chunks: Array<Buffer> = [];
+        stream
+          .on('data', (chunk: Buffer) => chunks.push(chunk))
+          .once('end', () => {
+            resolve(Buffer.concat(chunks));
+          })
+          .once('error', reject);
+      });
+    });
+  }
 }

src/devices/BaseDevice.tsx

@@ -145,4 +145,10 @@ export default class BaseDevice {
   archive(): any | null | undefined {
     return null;
   }
+
+  screenshot(): Promise<Buffer> {
+    return Promise.reject(
+      new Error('No screenshot support for current device'),
+    );
+  }
 }

src/devices/IOSDevice.tsx

@@ -10,6 +10,12 @@ import child_process from 'child_process';
 import BaseDevice from './BaseDevice';
 import JSONStream from 'JSONStream';
 import {Transform} from 'stream';
+import electron from 'electron';
+import fs from 'fs';
+import uuid from 'uuid/v1';
+import path from 'path';
+import {promisify} from 'util';
+import {exec} from 'child_process';
 
 type IOSLogLevel = 'Default' | 'Info' | 'Debug' | 'Error' | 'Fault';
 
@@ -42,6 +48,18 @@ export default class IOSDevice extends BaseDevice {
     this.log = this.startLogListener();
   }
 
+  screenshot(): Promise<Buffer> {
+    const tmpImageName = uuid() + '.png';
+    const tmpDirectory = (electron.app || electron.remote.app).getPath('temp');
+    const tmpFilePath = path.join(tmpDirectory, tmpImageName);
+    const command = `xcrun simctl io booted screenshot ${tmpFilePath}`;
+    return promisify(exec)(command)
+      .then(() => promisify(fs.readFile)(tmpFilePath))
+      .then(buffer => {
+        return promisify(fs.unlink)(tmpFilePath).then(() => buffer);
+      });
+  }
+
   teardown() {
     if (this.log) {
       this.log.kill();

src/index.js

@@ -25,6 +25,7 @@ export {default as constants} from './fb-stubs/constants.tsx';
 export {connect} from 'react-redux';
 export {selectPlugin} from './reducers/connections.tsx';
 export {getPluginKey, getPersistedState} from './utils/pluginUtils.tsx';
+export {writeBufferToFile, bufferToBlob} from './utils/screenshot.tsx';
 export type {Store, MiddlewareAPI} from './reducers/index.tsx';
 export {default as BaseDevice} from './devices/BaseDevice.tsx';
 

src/utils/screenshot.tsx

@@ -0,0 +1,31 @@
+/**
+ * Copyright 2018-present Facebook.
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ * @format
+ */
+
+import fs from 'fs';
+
+/**
+ * Writes a buffer to a specified file path.
+ * Returns a Promise which resolves to the file path.
+ */
+export const writeBufferToFile = (filePath: string, buffer: Buffer) => {
+  return new Promise((resolve, reject) => {
+    fs.writeFile(filePath, buffer, err => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve(filePath);
+      }
+    });
+  });
+};
+
+/**
+ * Creates a Blob from a Buffer
+ */
+export const bufferToBlob = (buffer: Buffer): Blob => {
+  return new Blob([buffer]);
+};