From 4771bea6d94aec72d164887019e6fbcf7f40b339 Mon Sep 17 00:00:00 2001
From: Rena-Yuan <84553146+Rena-Yuan@users.noreply.github.com>
Date: Mon, 28 Mar 2022 17:40:36 +0800
Subject: [PATCH] Add tidevice tool for screenshot. (#64)

* Add tidevice tool for screenshot.

* Modify the way to use tidevice

* Catch exception by tidevice-screenshot.
---
 README.md                  |  2 +-
 lyrebird_ios/apis.py       |  2 +-
 lyrebird_ios/ios_helper.py | 25 +++++++++++++++++++++----
 setup.py                   |  5 +++--
 4 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 290a7e3..fc47935 100644
--- a/README.md
+++ b/README.md
@@ -45,7 +45,7 @@
    ```bash
    ideviceinfo
    idevicescreenshot
-   idevicesinstaller -l
+   ideviceinstaller -l
    ```
 
 ### 安装
diff --git a/lyrebird_ios/apis.py b/lyrebird_ios/apis.py
index f035628..8e4b808 100644
--- a/lyrebird_ios/apis.py
+++ b/lyrebird_ios/apis.py
@@ -64,7 +64,7 @@ def take_screen_shot(device_id):
     device = device_service.devices.get(device_id)
     img_info = device.take_screen_shot()
     if img_info['returncode'] != 0:
-        return make_fail_response(img_info['result'].stdout.decode())
+        return make_fail_response(img_info['stdout'])
     timestamp = img_info.get('timestamp')
     return make_ok_response(imgUrl=f'/plugins/ios/api/src/screenshot/{device_id}?time={timestamp}')
 
diff --git a/lyrebird_ios/ios_helper.py b/lyrebird_ios/ios_helper.py
index bbebfb9..59c6e3a 100644
--- a/lyrebird_ios/ios_helper.py
+++ b/lyrebird_ios/ios_helper.py
@@ -5,6 +5,7 @@
 import plistlib
 import tempfile
 import subprocess
+import traceback
 import lyrebird
 from lyrebird.log import get_logger
 from . import wda_helper
@@ -271,13 +272,29 @@ def take_screen_shot(self):
         screen_shot_file = os.path.abspath(os.path.join(screenshot_dir, f'{file_name}_{timestamp}.png'))
         p = subprocess.run(f'{idevicescreenshot} -u {self.device_id} {screen_shot_file}',
                            shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        result = {
-            'returncode': p.returncode,
-            'result': p,
+        if p.returncode == 0:
+            return {
+                'returncode': p.returncode,
+                'stdout': p.stdout.decode(),
+                'screen_shot_file': screen_shot_file,
+                'timestamp': timestamp
+            }
+        import tidevice
+        try:
+            img = tidevice.Device(self.device_id).screenshot()
+        except AssertionError as e:
+            return_code = 1
+            stdout = f'Fail to get screenshot by tidevice.\n {str(e)}\n {traceback.format_exc()}'
+        else:
+            img.save(screen_shot_file)
+            return_code = 0
+            stdout = 'Success'
+        return {
+            'returncode': return_code,
+            'stdout': stdout,
             'screen_shot_file': screen_shot_file,
             'timestamp': timestamp
         }
-        return result
 
     def to_dict(self):
         device_info = {k: self.__dict__[k] for k in self.__dict__ if not k.startswith('_')}
diff --git a/setup.py b/setup.py
index 178b3c8..9602d23 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
 
 setup(
     name='lyrebird-ios',
-    version='0.3.8',
+    version='0.3.9',
     packages=['lyrebird_ios'],
     url='https://github.com/meituan/lyrebird-ios',
     author='HBQA',
@@ -29,7 +29,8 @@
     },
     install_requires=[
         'lyrebird',
-        'facebook-wda==0.8.1'
+        'facebook-wda==0.8.1',
+        'tidevice==0.6.12'
     ],
     extras_require={
         'dev': [