From afff7b0b645525f84d929355f7de7cbfe04e8e76 Mon Sep 17 00:00:00 2001
From: SnowLiu <309122873@qq.com>
Date: Sun, 11 Jan 2015 16:10:01 +0800
Subject: [PATCH] update screenshot

update screenshot
---
 appium/webdriver/mobilecommand.py | 3 +++
 appium/webdriver/webdriver.py     | 6 ++++++
 2 files changed, 9 insertions(+)

diff --git a/appium/webdriver/mobilecommand.py b/appium/webdriver/mobilecommand.py
index 6782ddae..45a8784e 100644
--- a/appium/webdriver/mobilecommand.py
+++ b/appium/webdriver/mobilecommand.py
@@ -54,3 +54,6 @@ class MobileCommand(object):
     START_ACTIVITY = 'startActivity'
     GET_SETTINGS = 'getSettings'
     UPDATE_SETTINGS = 'updateSettings'
+
+    #Custom
+    GET_SCREENSHOT = 'getScreenshot'
diff --git a/appium/webdriver/webdriver.py b/appium/webdriver/webdriver.py
index 20144379..f834e07c 100644
--- a/appium/webdriver/webdriver.py
+++ b/appium/webdriver/webdriver.py
@@ -678,6 +678,8 @@ def toggle_location_services(self):
         self.execute(Command.TOGGLE_LOCATION_SERVICES, {})
         return self
 
+    def screenshot(self):
+        return self.execute(Command.GET_SCREENSHOT)['value']
 
     def _addCommands(self):
         self.command_executor._commands[Command.CONTEXTS] = \
@@ -758,6 +760,10 @@ def _addCommands(self):
         self.command_executor._commands[Command.TOGGLE_LOCATION_SERVICES] = \
             ('POST', '/session/$sessionId/appium/device/toggle_location_services')
 
+        #custom screenshot
+        self.command_executor._commands[Command.GET_SCREENSHOT] = \
+            ('GET', '/session/$sessionId/screenshot')
+
 
 # monkeypatched method for WebElement
 def set_value(self, value):