From 769fac5855bef183ab4cdc5b634d3163b10d8e79 Mon Sep 17 00:00:00 2001
From: codeskyblue <codeskyblue@gmail.com>
Date: Thu, 9 May 2024 17:10:24 +0800
Subject: [PATCH] add push to dir support, close #101

---
 adbutils/_device.py             |  2 +-
 adbutils/errors.py              |  6 +++++-
 adbutils/sync.py                | 38 +++++++++++++++++++++++++++------
 test_real_device/conftest.py    |  9 ++++++++
 test_real_device/test_device.py | 17 ++++++++++++++-
 5 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/adbutils/_device.py b/adbutils/_device.py
index 9006801..ad55ffb 100644
--- a/adbutils/_device.py
+++ b/adbutils/_device.py
@@ -342,7 +342,7 @@ def framebuffer(self) -> Image.Image:
     @deprecated(deprecated_in="2.6.0", removed_in="3.0.0", current_version=__version__, details="use sync.push instead")
     def push(self, local: str, remote: str):
         """ alias for sync.push """
-        self.sync.push(local, remote)
+        return self.sync.push(local, remote)
 
     def create_connection(
         self, network: Network, address: Union[int, str]
diff --git a/adbutils/errors.py b/adbutils/errors.py
index ddf2665..f6c5872 100644
--- a/adbutils/errors.py
+++ b/adbutils/errors.py
@@ -35,4 +35,8 @@ def __init__(self, output: str):
         self.output = output
 
     def __str__(self):
-        return self.output
\ No newline at end of file
+        return self.output
+
+
+class AdbSyncError(AdbError):
+    """ sync error """
\ No newline at end of file
diff --git a/adbutils/sync.py b/adbutils/sync.py
index ad09b4a..9489369 100644
--- a/adbutils/sync.py
+++ b/adbutils/sync.py
@@ -5,6 +5,7 @@
 """
 
 
+import logging
 import struct
 import datetime
 import typing
@@ -17,8 +18,9 @@
 from adbutils._adb import BaseClient, AdbError
 from adbutils._proto import FileInfo
 from adbutils._utils import append_path
+from adbutils.errors import AdbSyncError
 
-
+logger = logging.getLogger(__name__)
 
 _OKAY = "OKAY"
 _FAIL = "FAIL"
@@ -80,10 +82,36 @@ def iter_directory(self, path: str):
     def list(self, path: str) -> typing.List[str]:
         return list(self.iter_directory(path))
 
-    def push(
+    def push(self, src: typing.Union[pathlib.Path, str, bytes, bytearray, typing.BinaryIO],
+             dst: typing.Union[pathlib.Path, str],
+             mode: int = 0o755,
+             check: bool = False) -> int:
+        """
+        Push file from local:src to device:dst
+
+        Args:
+            src: source file path
+            dst: destination file path or directory path
+            mode: file mode
+            check: check if push size is correct
+        
+        Returns:
+            total file size pushed
+        """
+        if isinstance(dst, pathlib.Path):
+            dst = dst.as_posix()
+        finfo = self.stat(dst)
+        if finfo.mode & stat.S_IFDIR != 0:
+            if not isinstance(src, (pathlib.Path, str)):
+                raise AdbSyncError("src should be a file path when dst is a directory")
+            dst = append_path(dst, pathlib.Path(src).name)
+            logger.debug("dst is a directory, update dst to %s", dst)
+        return self._push_file(src, dst, mode, check)
+    
+    def _push_file(
             self,
             src: typing.Union[pathlib.Path, str, bytes, bytearray, typing.BinaryIO],
-            dst: typing.Union[pathlib.Path, str],
+            dst: str,
             mode: int = 0o755,
             check: bool = False) -> int:  # yapf: disable
         # IFREG: File Regular
@@ -98,15 +126,13 @@ def push(
             if not hasattr(src, "read"):
                 raise TypeError("Invalid src type: %s" % type(src))
 
-        if isinstance(dst, pathlib.Path):
-            dst = dst.as_posix()
         path = dst + "," + str(stat.S_IFREG | mode)
         total_size = 0
         with self._prepare_sync(path, "SEND") as c:
             r = src if hasattr(src, "read") else open(src, "rb")
             try:
                 while True:
-                    chunk = r.read(4096)
+                    chunk = r.read(4096) # should not >64k
                     if not chunk:
                         mtime = int(datetime.datetime.now().timestamp())
                         c.conn.send(b"DONE" + struct.pack("<I", mtime))
diff --git a/test_real_device/conftest.py b/test_real_device/conftest.py
index 61988ad..a74abe6 100644
--- a/test_real_device/conftest.py
+++ b/test_real_device/conftest.py
@@ -16,6 +16,15 @@ def device_tmp_path(device: AdbDevice):
     yield tmp_path
     device.remove(tmp_path)
 
+
+@pytest.fixture
+def device_tmp_dir(device: AdbDevice):
+    tmp_path = "/data/local/tmp/adbutils-test"
+    device.shell("mkdir -p {}".format(tmp_path))
+    yield tmp_path
+    device.rmtree(tmp_path)
+
+
 @pytest.fixture
 def device_tmp_dir_path(device: AdbDevice):
     tmp_dir_path = "/sdcard/test_d"
diff --git a/test_real_device/test_device.py b/test_real_device/test_device.py
index 73ebe14..afe74a0 100644
--- a/test_real_device/test_device.py
+++ b/test_real_device/test_device.py
@@ -10,11 +10,13 @@
 import re
 import time
 import filecmp
+import uuid
 
 import pytest
 
 import adbutils
 from adbutils import AdbDevice, Network
+from adbutils.errors import AdbSyncError
 
 
 def test_shell(device: AdbDevice):
@@ -168,6 +170,17 @@ def test_sync_pull_file_push(device: AdbDevice, device_tmp_path, tmp_path: pathl
     assert b"Hello Android" == data
 
 
+def test_sync_push_to_dir(device: AdbDevice, device_tmp_dir, tmp_path: pathlib.Path):
+    random_data = str(uuid.uuid4()).encode()
+    src = io.BytesIO(random_data)
+    with pytest.raises(AdbSyncError):
+        device.sync.push(src, device_tmp_dir)
+    src_path = tmp_path.joinpath("random.txt")
+    src_path.write_bytes(random_data)
+    assert device.sync.push(src_path, device_tmp_dir) == len(random_data)
+    assert random_data == device.sync.read_bytes(device_tmp_dir + "/random.txt")
+
+
 def test_screenshot(device: AdbDevice):
     im = device.screenshot()
     assert im.mode == "RGB"
@@ -269,7 +282,9 @@ def are_dir_trees_equal(dir1, dir2):
     local_src_out_dir1 = tmp_path / 'dir1'
     local_src_out_dir2 = tmp_path / 'dir2'
 
-    device.push(local_src_in_dir, device_tmp_dir_path)
+    # TODO: push src support dir
+    # device.push(local_src_in_dir, device_tmp_dir_path)
+    device.adb_output("push", str(local_src_in_dir), device_tmp_dir_path)
 
     device.sync.pull_dir(device_tmp_dir_path, local_src_out_dir1)