Merge pull request #203 from drift-labs/sina/update-drift-user

Update Drift User and Drift Client

A bit of a stuffed PR, summary:
- Update Drift User the proper way to calculate `get_leverage` from the typescript SDK 
- Update `deposit` to use a temporary wSOL APA
- Update some constants to add the new listings

README.md

@@ -4,7 +4,8 @@
     <img src="docs/img/drift.png" width="30%" height="30%">
 </div>
 
-DriftPy is the Python client for the [Drift](https://www.drift.trade/) protocol. It allows you to trade and fetch data from Drift using Python.
+DriftPy is the Python client for the [Drift](https://www.drift.trade/) protocol.
+It allows you to trade and fetch data from Drift using Python.
 
 **[Read the full SDK documentation here!](https://drift-labs.github.io/v2-teacher/)**
 
@@ -16,15 +17,21 @@ pip install driftpy
 
 Note: requires Python >= 3.10.
 
-## ⚠️ IMPORTANT ⚠️
+
+## SDK Examples
+
+- `examples/` folder includes more examples of how to use the SDK including how to provide liquidity/become an lp, stake in the insurance fund, etc.
+
+
+## Note on using QuickNode
 
 If you are using QuickNode free plan, you *must* use `AccountSubscriptionConfig("demo")`, and you can only subscribe to 1 perp market and 1 spot market at a time.
 
 Non-QuickNode free RPCs (including the public mainnet-beta url) can use `cached` as well.
 
 Example setup for `AccountSubscriptionConfig("demo")`: 
 
-```
+```python
     # This example will listen to perp markets 0 & 1 and spot market 0
     # If you are listening to any perp markets, you must listen to spot market 0 or the SDK will break
 
@@ -48,31 +55,30 @@ If you intend to use `AccountSubscriptionConfig("demo)`, you *must* call `get_ma
 
 `get_markets_and_oracles` will return all the necessary `OracleInfo`s and `market_indexes` in order to use the SDK.
 
-## SDK Examples
-
-- `examples/` folder includes more examples of how to use the SDK including how to provide liquidity/become an lp, stake in the insurance fund, etc.
+# Development
 
 ## Setting Up Dev Env
 
 `bash setup.sh`
 
-# Development
 
 Ensure correct python version (using pyenv is recommended):
-```
+```bash
 pyenv install 3.10.11
 pyenv global 3.10.11
 poetry env use $(pyenv which python)
 ```
 
 Install dependencies:
-```
+```bash
 poetry install
 ```
 
-Run tests:
-```
-poetry run bash test-scripts/integration_test.sh
-poetry run bash test-scripts/math_tests.sh
-```
+To run tests, first ensure you have set up the RPC url, then run `pytest`:
+```bash
+export MAINNET_RPC_ENDPOINT="<YOUR_RPC_URL>"
+export DEVNET_RPC_ENDPOINT="https://api.devnet.solana.com" # or your own RPC
 
+poetry run pytest -v -s -x tests/ci/*.py
+poetry run pytest -v -s tests/math/*.py
+```

setup.sh

@@ -1,8 +1,3 @@
 git submodule update --init --recursive
-# build v2
 cd protocol-v2
-yarn && anchor build 
-# build dependencies for v2
-cd deps/serum-dex/dex && anchor build && cd ../../
-# go back to top-level
-cd ../
+yarn && anchor build

src/driftpy/constants/perp_markets.py

@@ -521,4 +521,11 @@ class PerpMarketConfig:
         oracle=Pubkey.from_string("DpJz7rjTJLxxnuqrqZTUjMWtnaMFAEfZUv5ATdb9HTh1"),
         oracle_source=OracleSource.Prelaunch(),
     ),
+    PerpMarketConfig(
+        symbol="MOTHER-PERP",
+        base_asset_symbol="MOTHER",
+        market_index=44,
+        oracle=Pubkey.from_string("56ap2coZG7FPWUigVm9XrpQs3xuCwnwQaWtjWZcffEUG"),
+        oracle_source=OracleSource.PythPull(),
+    ),
 ]

src/driftpy/drift_client.py

@@ -1,57 +1,69 @@
 import json
 import os
-import anchorpy
-from deprecated import deprecated
-import requests
-from solders.pubkey import Pubkey
-from solders.keypair import Keypair
-from solders.transaction import TransactionVersion, Legacy
-from solders.instruction import Instruction
-from solders.system_program import ID
-from solders.sysvar import RENT
-from solders.signature import Signature
-from solders.address_lookup_table_account import AddressLookupTableAccount
-from solana.rpc.async_api import AsyncClient
-from solana.rpc.types import TxOpts
-from solana.rpc.commitment import Processed
-from solana.transaction import AccountMeta
-from solders.compute_budget import set_compute_unit_limit, set_compute_unit_price
-from spl.token.constants import TOKEN_PROGRAM_ID
-from spl.token.instructions import get_associated_token_address
-from anchorpy import Program, Context, Idl, Provider, Wallet
 from pathlib import Path
+import random
+import string
+from typing import List, Optional, Tuple, Union
 
+import anchorpy
+from anchorpy import Context
+from anchorpy import Idl
+from anchorpy import Program
+from anchorpy import Provider
+from anchorpy import Wallet
+from deprecated import deprecated
 import driftpy
 from driftpy.account_subscription_config import AccountSubscriptionConfig
+from driftpy.accounts import *
 from driftpy.address_lookup_table import get_address_lookup_table
-from driftpy.constants import BASE_PRECISION, PRICE_PRECISION
-from driftpy.constants.numeric_constants import (
-    QUOTE_SPOT_MARKET_INDEX,
-)
+from driftpy.addresses import get_sequencer_public_key_and_bump
+from driftpy.constants import BASE_PRECISION
+from driftpy.constants import PRICE_PRECISION
+from driftpy.constants.config import configs
+from driftpy.constants.config import decode_account
+from driftpy.constants.config import DEVNET_SEQUENCER_PROGRAM_ID
+from driftpy.constants.config import DRIFT_PROGRAM_ID
+from driftpy.constants.config import DriftEnv
+from driftpy.constants.config import SEQUENCER_PROGRAM_ID
+from driftpy.constants.numeric_constants import QUOTE_SPOT_MARKET_INDEX
+from driftpy.constants.spot_markets import WRAPPED_SOL_MINT
 from driftpy.decode.utils import decode_name
 from driftpy.drift_user import DriftUser
-from driftpy.accounts import *
-
-from driftpy.constants.config import (
-    DriftEnv,
-    DRIFT_PROGRAM_ID,
-    configs,
-    SEQUENCER_PROGRAM_ID,
-    DEVNET_SEQUENCER_PROGRAM_ID,
-    decode_account,
-)
-
-from typing import Tuple, Union, Optional, List
-from driftpy.drift_user_stats import DriftUserStats, UserStatsSubscriptionConfig
+from driftpy.drift_user_stats import DriftUserStats
+from driftpy.drift_user_stats import UserStatsSubscriptionConfig
 from driftpy.math.perp_position import is_available
-from driftpy.math.spot_position import is_spot_position_available
 from driftpy.math.spot_market import cast_to_spot_precision
+from driftpy.math.spot_position import is_spot_position_available
 from driftpy.name import encode_name
 from driftpy.tx.standard_tx_sender import StandardTxSender
-from driftpy.tx.types import TxSender, TxSigAndSlot
-from driftpy.addresses import get_sequencer_public_key_and_bump
-from spl.token.constants import ASSOCIATED_TOKEN_PROGRAM_ID
+from driftpy.tx.types import TxSender
+from driftpy.tx.types import TxSigAndSlot
+import requests
+from solana.rpc.async_api import AsyncClient
+from solana.rpc.commitment import Processed
+from solana.rpc.types import TxOpts
+from solana.transaction import AccountMeta
+from solders import system_program
+from solders.address_lookup_table_account import AddressLookupTableAccount
+from solders.compute_budget import set_compute_unit_limit
+from solders.compute_budget import set_compute_unit_price
+from solders.instruction import Instruction
+from solders.keypair import Keypair
+from solders.pubkey import Pubkey
+from solders.signature import Signature
+from solders.system_program import ID
 from solders.system_program import ID as SYS_PROGRAM_ID
+from solders.sysvar import RENT
+from solders.transaction import Legacy
+from solders.transaction import TransactionVersion
+from spl.token.constants import ASSOCIATED_TOKEN_PROGRAM_ID
+from spl.token.constants import TOKEN_PROGRAM_ID
+from spl.token.instructions import close_account
+from spl.token.instructions import CloseAccountParams
+from spl.token.instructions import get_associated_token_address
+from spl.token.instructions import initialize_account
+from spl.token.instructions import InitializeAccountParams
+
 
 DEFAULT_USER_NAME = "Main Account"
 
@@ -863,12 +875,52 @@ def get_initialize_user_instructions(
         )
         return initialize_user_account_ix
 
+    def random_string(self, length: int) -> str:
+        return "".join(random.choices(string.ascii_letters + string.digits, k=length))
+
+    async def get_wrapped_sol_account_creation_ixs(
+        self, amount: int, include_rent: bool = True
+    ) -> (List[Instruction], Pubkey):
+        wallet_pubkey = self.wallet.public_key
+        seed = self.random_string(32)
+        wrapped_sol_account = Pubkey.create_with_seed(
+            wallet_pubkey, seed, TOKEN_PROGRAM_ID
+        )
+        result = {"ixs": [], "pubkey": wrapped_sol_account}
+
+        LAMPORTS_PER_SOL: int = 1_000_000_000
+        rent_space_lamports = int(LAMPORTS_PER_SOL / 100)
+        lamports = amount + rent_space_lamports if include_rent else rent_space_lamports
+
+        create_params = system_program.CreateAccountWithSeedParams(
+            from_pubkey=wallet_pubkey,
+            to_pubkey=wrapped_sol_account,
+            base=wallet_pubkey,
+            seed=seed,
+            lamports=lamports,
+            space=165,
+            owner=TOKEN_PROGRAM_ID,
+        )
+
+        result["ixs"].append(system_program.create_account_with_seed(create_params))
+
+        initialize_params = InitializeAccountParams(
+            program_id=TOKEN_PROGRAM_ID,
+            account=wrapped_sol_account,
+            mint=WRAPPED_SOL_MINT,
+            owner=wallet_pubkey,
+        )
+
+        result["ixs"].append(initialize_account(initialize_params))
+
+        return result["ixs"], result["pubkey"]
+
     async def deposit(
         self,
         amount: int,
         spot_market_index: int,
-        user_token_account: Pubkey = None,
-        sub_account_id: int = None,
+        user_token_account: Pubkey,
+        sub_account_id: Optional[int] = None,
         reduce_only=False,
         user_initialized=True,
     ):
@@ -886,30 +938,39 @@ async def deposit(
             str: sig
         """
         tx_sig_and_slot = await self.send_ixs(
-            [
-                self.get_deposit_collateral_ix(
-                    amount,
-                    spot_market_index,
-                    user_token_account,
-                    sub_account_id,
-                    reduce_only,
-                    user_initialized,
-                )
-            ]
+            await self.get_deposit_collateral_ix(
+                amount,
+                spot_market_index,
+                user_token_account,
+                sub_account_id,
+                reduce_only,
+                user_initialized,
+            )
         )
         self.last_spot_market_seen_cache[spot_market_index] = tx_sig_and_slot.slot
         return tx_sig_and_slot
 
-    def get_deposit_collateral_ix(
+    async def get_deposit_collateral_ix(
         self,
         amount: int,
         spot_market_index: int,
-        user_token_account: Pubkey = None,
-        sub_account_id: int = None,
-        reduce_only=False,
-        user_initialized=True,
-    ) -> Instruction:
+        user_token_account: Pubkey,
+        sub_account_id: Optional[int] = None,
+        reduce_only: Optional[bool] = False,
+        user_initialized: Optional[bool] = True,
+    ) -> List[Instruction]:
+
         sub_account_id = self.get_sub_account_id_for_ix(sub_account_id)
+        spot_market_account = self.get_spot_market_account(spot_market_index)
+        if not spot_market_account:
+            raise Exception("Spot market account not found")
+
+        is_sol_market = spot_market_account.mint == WRAPPED_SOL_MINT
+        signer_authority = self.wallet.public_key
+
+        create_WSOL_token_account = (
+            is_sol_market and user_token_account == signer_authority
+        )
 
         if user_initialized:
             remaining_accounts = self.get_remaining_accounts(
@@ -919,6 +980,13 @@ def get_deposit_collateral_ix(
         else:
             raise Exception("not implemented...")
 
+        instructions = []
+
+        if create_WSOL_token_account:
+            ixs, ata_pubkey = await self.get_wrapped_sol_account_creation_ixs(amount)
+            instructions.extend(ixs)
+            user_token_account = ata_pubkey
+
         user_token_account = (
             user_token_account
             if user_token_account is not None
@@ -932,7 +1000,7 @@ def get_deposit_collateral_ix(
         user_account_public_key = get_user_account_public_key(
             self.program_id, self.authority, sub_account_id
         )
-        return self.program.instruction["deposit"](
+        deposit_ix = self.program.instruction["deposit"](
             spot_market_index,
             amount,
             reduce_only,
@@ -950,6 +1018,16 @@ def get_deposit_collateral_ix(
                 remaining_accounts=remaining_accounts,
             ),
         )
+        instructions.append(deposit_ix)
+        close_account_params = CloseAccountParams(
+            program_id=TOKEN_PROGRAM_ID,
+            account=ata_pubkey,
+            dest=signer_authority,
+            owner=signer_authority,
+        )
+        close_account_ix = close_account(close_account_params)
+        instructions.append(close_account_ix)
+        return instructions
 
     async def withdraw(
         self,