uniswap-python · ErikBjare · Jul 14, 2022 · Jul 14, 2022 · Jul 14, 2022 · Jul 14, 2022
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,6 +27,7 @@ python = "^3.7.2"
 web3 = "^5.23.0"
 click = "^8.0.3"
 python-dotenv = "*"
+typing-extensions = "^4.3.0"
 
 [tool.poetry.dev-dependencies]
 mypy = "*"

diff --git a/uniswap/decorators.py b/uniswap/decorators.py
@@ -1,20 +1,28 @@
 import functools
-from typing import Callable, Any, List, Dict, TYPE_CHECKING
+from typing import Callable, List, TYPE_CHECKING, TypeVar, Optional
+from typing_extensions import ParamSpec, Concatenate
 
+from .types import AddressLike
 from .constants import ETH_ADDRESS
 
 if TYPE_CHECKING:
     from .uniswap import Uniswap
 
 
-def check_approval(method: Callable) -> Callable:
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def check_approval(
+    method: Callable[Concatenate["Uniswap", P], T]
+) -> Callable[Concatenate["Uniswap", P], T]:
     """Decorator to check if user is approved for a token. It approves them if they
     need to be approved."""
 
     @functools.wraps(method)
-    def approved(self: Any, *args: Any, **kwargs: Any) -> Any:
+    def approved(self: "Uniswap", *args: P.args, **kwargs: P.kwargs) -> T:
         # Check to see if the first token is actually ETH
-        token = args[0] if args[0] != ETH_ADDRESS else None
+        token: Optional[AddressLike] = args[0] if args[0] != ETH_ADDRESS else None  # type: ignore
         token_two = None
 
         # Check second token, if needed
@@ -32,8 +40,14 @@ def approved(self: Any, *args: Any, **kwargs: Any) -> Any:
     return approved
 
 
-def supports(versions: List[int]) -> Callable:
-    def g(f: Callable) -> Callable:
+def supports(
+    versions: List[int],
+) -> Callable[
+    [Callable[Concatenate["Uniswap", P], T]], Callable[Concatenate["Uniswap", P], T]
+]:
+    def g(
+        f: Callable[Concatenate["Uniswap", P], T]
+    ) -> Callable[Concatenate["Uniswap", P], T]:
         if f.__doc__ is None:
             f.__doc__ = ""
         f.__doc__ += """\n\n
@@ -43,7 +57,7 @@ def g(f: Callable) -> Callable:
         )
 
         @functools.wraps(f)
-        def check_version(self: "Uniswap", *args: List, **kwargs: Dict) -> Any:
+        def check_version(self: "Uniswap", *args: P.args, **kwargs: P.kwargs) -> T:
             if self.version not in versions:
                 raise Exception(
                     f"Function {f.__name__} does not support version {self.version} of Uniswap passed to constructor"

diff --git a/uniswap/types.py b/uniswap/types.py
@@ -1,6 +1,5 @@
 from typing import Union
-from web3.eth import Contract  # noqa: F401
-from web3.types import Address, ChecksumAddress
+from eth_typing.evm import Address, ChecksumAddress
 
 
 AddressLike = Union[Address, ChecksumAddress]
diff --git a/uniswap/uniswap.py b/uniswap/uniswap.py
@@ -316,7 +316,7 @@ def _get_eth_token_output_price(
             route = [self.get_weth_address(), token]
             price = self.router.functions.getAmountsIn(qty, route).call()[0]
         elif self.version == 3:
-            if not fee:
+            if fee is None:
                 logger.warning("No fee set, assuming 0.3%")
                 fee = 3000
             price = Wei(
@@ -1256,17 +1256,24 @@ def get_weth_address(self) -> ChecksumAddress:
             address: ChecksumAddress = self.router.functions.WETH().call()
         elif self.version == 3:
             address = self.router.functions.WETH9().call()
+        else:
+            raise ValueError  # pragma: no cover
         return address
 
     @supports([2, 3])
     def get_raw_price(
-        self, token_in: AddressLike, token_out: AddressLike, fee: int = 3000
+        self, token_in: AddressLike, token_out: AddressLike, fee: int = None
     ) -> float:
         """
         Returns current price for pair of tokens [token_in, token_out] regrading liquidity that is being locked in the pool
         Parameter `fee` is required for V3 only, can be omitted for V2
         Requires pair [token_in, token_out] having direct pool
         """
+        if not fee:
+            fee = 3000
+            if self.version == 3:
+                logger.warning("No fee set, assuming 0.3%")
+
         if token_in == ETH_ADDRESS:
             token_in = self.get_weth_address()
         if token_out == ETH_ADDRESS:

diff --git a/uniswap/util.py b/uniswap/util.py
@@ -5,8 +5,9 @@
 
 from web3 import Web3
 from web3.exceptions import NameNotFound
+from web3.contract import Contract
 
-from .types import AddressLike, Address, Contract
+from .types import AddressLike, Address
 
 
 def _str_to_addr(s: Union[AddressLike, str]) -> Address: