Merge pull request #194 from openatx/copilot/fix-193

Optimize shell and shell2 method type hints using @overload and Literal

Author: codeskyblue

adbutils/_device_base.py

@@ -11,7 +11,12 @@
 import subprocess
 import threading
 import typing
-from typing import List, Optional, Union
+from typing import List, Optional, Union, overload
+
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
 
 from PIL import Image, UnidentifiedImageError
 from deprecation import deprecated
@@ -172,6 +177,47 @@ def open_shell(self, cmdargs: Union[str, list, tuple]) -> AdbConnection:
         c.check_okay()
         return c
 
+    @overload
+    def shell(
+        self,
+        cmdargs: Union[str, list, tuple],
+        stream: Literal[True],
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        encoding: str | None = "utf-8",
+        rstrip=True,
+    ) -> AdbConnection: ...
+    
+    @overload
+    def shell(
+        self,
+        cmdargs: Union[str, list, tuple],
+        *,
+        stream: Literal[False] = False,
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        encoding: Literal[None],
+        rstrip=True,
+    ) -> bytes: ...
+    
+    @overload
+    def shell(
+        self,
+        cmdargs: Union[str, list, tuple],
+        stream: Literal[False] = False,
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        *,
+        encoding: str,
+        rstrip=True,
+    ) -> str: ...
+    
+    @overload
+    def shell(
+        self,
+        cmdargs: Union[str, list, tuple],
+        stream: Literal[False] = False,
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        rstrip=True,
+    ) -> str: ...
+
     def shell(
         self,
         cmdargs: Union[str, list, tuple],
@@ -218,6 +264,37 @@ def shell(
             return output.rstrip() if rstrip else output
         return output
 
+    @overload
+    def shell2(
+        self,
+        cmdargs: Union[str, list, tuple],
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        *,
+        encoding: Literal[None],
+        rstrip=False,
+        v2=False,
+    ) -> ShellReturnRaw: ...
+    
+    @overload
+    def shell2(
+        self,
+        cmdargs: Union[str, list, tuple],
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        *,
+        encoding: str,
+        rstrip=False,
+        v2=False,
+    ) -> ShellReturn: ...
+    
+    @overload
+    def shell2(
+        self,
+        cmdargs: Union[str, list, tuple],
+        timeout: Optional[float] = _DEFAULT_SOCKET_TIMEOUT,
+        rstrip=False,
+        v2=False,
+    ) -> ShellReturn: ...
+
     def shell2(
         self,
         cmdargs: Union[str, list, tuple],

adbutils/_interfaces.py

@@ -1,9 +1,14 @@
 import abc
-from typing import List, Optional, Union
+from typing import List, Optional, Union, overload
+
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
 
 from adbutils._adb import AdbConnection, Network
 from adbutils.sync import Sync
-from adbutils._proto import ShellReturn
+from adbutils._proto import ShellReturn, ShellReturnRaw
 
 class AbstractShellDevice(abc.ABC):
     @abc.abstractmethod
@@ -17,10 +22,26 @@ def sync(self) -> Sync:
 
 
 class AbstractDevice(abc.ABC):
+    @overload
+    @abc.abstractmethod
+    def shell(self, cmd: str, stream: Literal[True]) -> AdbConnection: ...
+    
+    @overload
+    @abc.abstractmethod
+    def shell(self, cmd: str, stream: Literal[False] = False) -> str: ...
+    
     @abc.abstractmethod
     def shell(self, cmd: str, stream: bool = False) -> Union[str, AdbConnection]:
         pass
 
+    @overload
+    @abc.abstractmethod
+    def shell2(self, cmd: str, encoding: Literal[None]) -> ShellReturnRaw: ...
+    
+    @overload
+    @abc.abstractmethod
+    def shell2(self, cmd: str, encoding: str = "utf-8") -> ShellReturn: ...
+    
     @abc.abstractmethod
     def shell2(self, cmd: str) -> ShellReturn:
         pass