From 8a99c5d6b0fb58bf259fda4bb5830e03ef07a370 Mon Sep 17 00:00:00 2001
From: kayibal <alan@datarevenue.com>
Date: Fri, 25 Jun 2021 10:33:37 +0200
Subject: [PATCH] Support swapping directly with WETH

---
 tests/test_uniswap.py | 27 +++++++++++++++++++++++++++
 uniswap/uniswap.py    | 23 ++++++++++++++++++++---
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/tests/test_uniswap.py b/tests/test_uniswap.py
index 8a75ca0..5b87c00 100644
--- a/tests/test_uniswap.py
+++ b/tests/test_uniswap.py
@@ -330,3 +330,30 @@ def test_make_trade_output(
             balance_after = client.get_token_balance(output_token)
             if output_token != self.eth:
                 assert balance_before + qty == balance_after
+
+    @pytest.mark.parametrize(
+        "input_token, output_token, route, expectation",
+        [
+            (dai, weth, [dai, weth], does_not_raise),
+            (weth, dai, [weth, dai], does_not_raise),
+            (usdc, dai, [usdc, weth, dai], does_not_raise),
+            (weth, eth, [], lambda: pytest.raises(ValueError)),
+            (eth, weth, [], lambda: pytest.raises(ValueError)),
+        ],
+    )
+    def test_get_v2_trade_route(
+        self,
+        input_token,
+        output_token,
+        route,
+        expectation,
+        client
+    ):
+        if client.version == 1:
+            pytest.skip(
+                "Not supported in this version of Uniswap, or at least no liquidity"
+            )
+        with expectation():
+            res = client._get_v2_trade_route(input_token, output_token)
+
+            assert res == route
diff --git a/uniswap/uniswap.py b/uniswap/uniswap.py
index 9242bce..041c61a 100644
--- a/uniswap/uniswap.py
+++ b/uniswap/uniswap.py
@@ -330,7 +330,7 @@ def _get_token_token_output_price(
                 elif is_same_address(token1, self.get_weth_address()):
                     return int(self._get_token_eth_output_price(token0, Wei(qty), fee))
 
-                route = [token0, self.get_weth_address(), token1]
+                route = self._get_v2_trade_route(token0, token1)
                 logger.warning(f"No route specified, assuming route: {route}")
 
         if self.version == 2:
@@ -558,11 +558,12 @@ def _token_to_token_swap_input(
                     input_token, output_token, qty, fee=fee
                 )
             )
+            route = self._get_v2_trade_route(input_token, output_token)
             return self._build_and_send_tx(
                 self.router.functions.swapExactTokensForTokens(
                     qty,
                     min_tokens_bought,
-                    [input_token, self.get_weth_address(), output_token],
+                    route,
                     recipient,
                     self._deadline(),
                 ),
@@ -595,6 +596,22 @@ def _token_to_token_swap_input(
         else:
             raise ValueError
 
+    def _get_v2_trade_route(
+        self, input_token: AddressLike, output_token: AddressLike
+    ) -> List[AddressLike]:
+        """Gets trade route simply through WETH.
+
+        Deals with two special cases when trading directly with WETH.
+        """
+        weth = self.get_weth_address()
+        if {weth, ETH_ADDRESS} == {input_token, output_token}:
+            raise ValueError("Swapping between ETH and WETH not supported!")
+        if input_token == weth:
+            return [weth, output_token]
+        elif output_token == weth:
+            return [input_token, weth]
+        return [input_token, weth, output_token]
+
     def _eth_to_token_swap_output(
         self,
         output_token: AddressLike,
@@ -736,7 +753,7 @@ def _token_to_token_swap_output(
                 self.router.functions.swapTokensForExactTokens(
                     qty,
                     amount_in_max,
-                    [input_token, self.get_weth_address(), output_token],
+                    self._get_v2_trade_route(input_token, output_token),
                     recipient,
                     self._deadline(),
                 ),