more reliable screen capture

droidbot/device.py

@@ -177,6 +177,9 @@ def disconnect(self):
             self.view_client.disconnect()
         self.logcat.terminate()
         self.state_monitor.stop()
+        temp_dir = os.path.join(self.output_dir, "temp")
+        if os.path.exists(temp_dir):
+            os.rmdir(temp_dir)
 
     def get_telnet(self):
         """
@@ -643,18 +646,32 @@ def pull_file(self, remote_file, local_file):
         self.get_adb().run_cmd(["pull", remote_file, local_file])
 
     def take_screenshot(self):
-        image = None
-
-        received = self.get_adb().shell("screencap -p").replace("\r\n", "\n")
-        import StringIO
-        stream = StringIO.StringIO(received)
-
-        try:
-            from PIL import Image
-            image = Image.open(stream)
-        except IOError as e:
-            self.logger.warning("exception in take_screenshot: %s" % e)
-        return image
+        # image = None
+        #
+        # received = self.get_adb().shell("screencap -p").replace("\r\n", "\n")
+        # import StringIO
+        # stream = StringIO.StringIO(received)
+        #
+        # try:
+        #     from PIL import Image
+        #     image = Image.open(stream)
+        # except IOError as e:
+        #     self.logger.warning("exception in take_screenshot: %s" % e)
+        # return image
+        from datetime import datetime
+        tag = datetime.now().strftime("%Y-%m-%d_%H%M%S")
+
+        remote_image_path = "/sdcard/screen_%s.png" % tag
+        local_image_dir = os.path.join(self.output_dir, "temp")
+        if not os.path.exists(local_image_dir):
+            os.mkdir(local_image_dir)
+        local_image_path = os.path.join(local_image_dir, "screen_%s.png" % tag)
+
+        self.get_adb().shell("screencap -p %s" % remote_image_path)
+        self.pull_file(remote_image_path, local_image_path)
+        self.get_adb().shell("rm %s" % remote_image_path)
+
+        return local_image_path
 
     def get_current_state(self):
         self.logger.info("getting current device state...")
@@ -663,13 +680,13 @@ def get_current_state(self):
             view_client_views = self.dump_views()
             foreground_activity = self.get_top_activity_name()
             background_services = self.get_service_names()
-            screenshot = self.take_screenshot()
+            screenshot_path = self.take_screenshot()
             self.logger.info("finish getting current device state...")
             current_state = DeviceState(self,
                                         view_client_views=view_client_views,
                                         foreground_activity=foreground_activity,
                                         background_services=background_services,
-                                        screenshot=screenshot)
+                                        screenshot_path=screenshot_path)
         except Exception as e:
             self.logger.warning("exception in get_current_state: %s" % e)
             import traceback
@@ -711,7 +728,7 @@ class DeviceState(object):
     the state of the current device
     """
 
-    def __init__(self, device, view_client_views, foreground_activity, background_services, tag=None, screenshot=None):
+    def __init__(self, device, view_client_views, foreground_activity, background_services, tag=None, screenshot_path=None):
         self.device = device
         self.view_client_views = view_client_views
         self.foreground_activity = foreground_activity
@@ -720,7 +737,7 @@ def __init__(self, device, view_client_views, foreground_activity, background_se
             from datetime import datetime
             tag = datetime.now().strftime("%Y-%m-%d_%H%M%S")
         self.tag = tag
-        self.screenshot = screenshot
+        self.screenshot_path = screenshot_path
         self.views = self.views2list(view_client_views)
         self.view_str = self.get_state_str()
 
@@ -784,13 +801,14 @@ def save2dir(self, output_dir=None):
             if not os.path.exists(output_dir):
                 os.mkdir(output_dir)
             state_json_file_path = "%s/state_%s.json" % (output_dir, self.tag)
-            screenshot_file_path = "%s/screenshot_%s.png" % (output_dir, self.tag)
+            screenshot_output_path = "%s/screenshot_%s.png" % (output_dir, self.tag)
             state_json_file = open(state_json_file_path, "w")
             state_json_file.write(self.to_json())
             state_json_file.close()
-            from PIL.Image import Image
-            if isinstance(self.screenshot, Image):
-                self.screenshot.save(screenshot_file_path)
+            subprocess.check_call(["cp", self.screenshot_path, screenshot_output_path])
+            # from PIL.Image import Image
+            # if isinstance(self.screenshot_path, Image):
+            #     self.screenshot_path.save(screenshot_output_path)
         except Exception as e:
             self.device.logger.warning("saving state to dir failed: " + e.message)