black

drift-labs · Jul 22, 2024 · c1cf7ec · c1cf7ec
1 parent 3cb96ae
commit c1cf7ec
Show file tree

Hide file tree

Showing 6 changed files with 716 additions and 135 deletions.
diff --git a/examples/place_and_take.py b/examples/place_and_take.py
@@ -1,4 +1,4 @@
-'''
+"""
 This example demonstrates use of the `place_and_take` instruction for perp and spot markets.
 1) get L3 dlob from dlob server
 2) build MakerInfo[] using UserMap
@@ -19,7 +19,8 @@
 
 Usage:
     python3 examples/place_and_take.py
-'''
+"""
+
 import os
 import sys
 import asyncio
@@ -38,60 +39,82 @@
 
 from solana.rpc.async_api import AsyncClient
 
-from driftpy.types import MakerInfo, MarketType, OrderType, OrderParams, PositionDirection, TxParams, is_variant
+from driftpy.types import (
+    MakerInfo,
+    MarketType,
+    OrderType,
+    OrderParams,
+    PositionDirection,
+    TxParams,
+    is_variant,
+)
 from driftpy.account_subscription_config import AccountSubscriptionConfig
 from driftpy.drift_client import DriftClient
 from driftpy.constants.numeric_constants import BASE_PRECISION, PRICE_PRECISION
 
 import requests
 
+
 def get_l3(market_type: str, market_index: int):
     url = f"https://dlob.drift.trade/l3?marketType={market_type}&marketIndex={market_index}"
     response = requests.get(url)
     return response.json()
 
 
-async def demo_perp_place_and_take(drift_client: DriftClient, perp_market_index: int, perp_trade_direction: PositionDirection, perp_base_amount: float):
+async def demo_perp_place_and_take(
+    drift_client: DriftClient,
+    perp_market_index: int,
+    perp_trade_direction: PositionDirection,
+    perp_base_amount: float,
+):
     user_account = drift_client.get_user_account()
     user_stats = drift_client.get_user_stats()
     print(f"user referrer info: {user_stats.get_referrer_info()}")
     print(f"Placing orders under authority:   {user_account.authority}")
-    print(f"Placing orders under userAccount: {drift_client.get_user_account_public_key()}, subaccount_id: {user_account.sub_account_id}")
+    print(
+        f"Placing orders under userAccount: {drift_client.get_user_account_public_key()}, subaccount_id: {user_account.sub_account_id}"
+    )
 
     # user map stores User account info, this is needed to get the UserStats account for makers we want to take against
-    user_map = UserMap(UserMapConfig(
-        drift_client,
-        WebsocketConfig(),
-        drift_client.connection,
-        skip_initial_load=True, # will lazy load if True
-    ))
+    user_map = UserMap(
+        UserMapConfig(
+            drift_client,
+            WebsocketConfig(),
+            drift_client.connection,
+            skip_initial_load=True,  # will lazy load if True
+        )
+    )
     await user_map.subscribe()
 
     oracle_price = drift_client.get_oracle_price_data_for_perp_market(0)
     if oracle_price is None:
         raise Exception("Failed to get oracle price")
     print(f"Perp market oracle price: {oracle_price.price / PRICE_PRECISION}")
 
-    book_side_to_take = 'bids'
-    if is_variant(perp_trade_direction, 'long'):
-        book_side_to_take = 'asks'
+    book_side_to_take = "bids"
+    if is_variant(perp_trade_direction, "long"):
+        book_side_to_take = "asks"
 
     # demonstrating passing in 3 makers, the contract will give the best price available at fill time
     # build maker info map
     perp_l3 = get_l3("perp", perp_market_index)
     maker_infos = []
     book_side = perp_l3[book_side_to_take]
     num_makers = min(3, len(book_side))
-    for maker in [level['maker'] for level in book_side[:num_makers]]:
+    for maker in [level["maker"] for level in book_side[:num_makers]]:
         maker_pubkey = Pubkey.from_string(maker)
         maker_user = await user_map.must_get(maker)
         maker_user_account = maker_user.get_user_account()
-        maker_infos.append(MakerInfo(
-            maker=maker_pubkey,
-            maker_stats=get_user_stats_account_public_key(drift_client.program_id, maker_user_account.authority),
-            maker_user_account=maker_user_account,
-            order=None,
-        ))
+        maker_infos.append(
+            MakerInfo(
+                maker=maker_pubkey,
+                maker_stats=get_user_stats_account_public_key(
+                    drift_client.program_id, maker_user_account.authority
+                ),
+                maker_user_account=maker_user_account,
+                order=None,
+            )
+        )
 
     print(f"simulating place_and_take_perp tx with {len(maker_infos)} makers")
     ixs = [
@@ -107,52 +130,65 @@ async def demo_perp_place_and_take(drift_client: DriftClient, perp_market_index:
                 price=oracle_price.price,
             ),
             maker_infos,
-            referrer_info=user_stats.get_referrer_info(), # if your UserAccount was referred, this is required.
-        )
+            referrer_info=user_stats.get_referrer_info(),  # if your UserAccount was referred, this is required.
+        ),
     ]
     tx = await drift_client.tx_sender.get_versioned_tx(
-        ixs,
-        drift_client.wallet.payer,
-        [await drift_client.fetch_market_lookup_table()])
-    tx_sim = await drift_client.connection.simulate_transaction(tx) # simulate because it's been hard to land a tx
+        ixs, drift_client.wallet.payer, [await drift_client.fetch_market_lookup_table()]
+    )
+    tx_sim = await drift_client.connection.simulate_transaction(
+        tx
+    )  # simulate because it's been hard to land a tx
     print(f"Error: {tx_sim.value.err}")
     print(f"CU used: {tx_sim.value.units_consumed}")
     print(f"logs:")
     [print(log) for log in tx_sim.value.logs]
 
 
-async def demo_spot_place_and_take(drift_client: DriftClient, user_map: UserMap, spot_market_index: int, spot_trade_direction: PositionDirection, spot_base_amount: float):
+async def demo_spot_place_and_take(
+    drift_client: DriftClient,
+    user_map: UserMap,
+    spot_market_index: int,
+    spot_trade_direction: PositionDirection,
+    spot_base_amount: float,
+):
     user_account = drift_client.get_user_account()
     user_stats = drift_client.get_user_stats()
     print(f"user referrer info: {user_stats.get_referrer_info()}")
     print(f"Placing orders under authority:   {user_account.authority}")
-    print(f"Placing orders under userAccount: {drift_client.get_user_account_public_key()}, subaccount_id: {user_account.sub_account_id}")
+    print(
+        f"Placing orders under userAccount: {drift_client.get_user_account_public_key()}, subaccount_id: {user_account.sub_account_id}"
+    )
 
     oracle_price = drift_client.get_oracle_price_data_for_spot_market(spot_market_index)
     if oracle_price is None:
         raise Exception("Failed to get oracle price")
     print(f"Spot market oracle price: {oracle_price.price / PRICE_PRECISION}")
 
-    book_side_to_take = 'bids'
-    if is_variant(spot_trade_direction, 'long'):
-        book_side_to_take = 'asks'
+    book_side_to_take = "bids"
+    if is_variant(spot_trade_direction, "long"):
+        book_side_to_take = "asks"
 
     # demonstrating passing in 3 makers, the contract will give the best price available at fill time
     # build maker info map
     spot_l3 = get_l3("spot", spot_market_index)
     maker_infos = []
     book_side = spot_l3[book_side_to_take]
     num_makers = min(3, len(book_side))
-    for maker in [level['maker'] for level in book_side[:num_makers]]:
+    for maker in [level["maker"] for level in book_side[:num_makers]]:
         maker_pubkey = Pubkey.from_string(maker)
         maker_user = await user_map.must_get(maker)
         maker_user_account = maker_user.get_user_account()
-        maker_infos.append(MakerInfo(
-            maker=maker_pubkey,
-            maker_stats=get_user_stats_account_public_key(drift_client.program_id, maker_user_account.authority),
-            maker_user_account=maker_user_account,
-            order=None,
-        ))
+        maker_infos.append(
+            MakerInfo(
+                maker=maker_pubkey,
+                maker_stats=get_user_stats_account_public_key(
+                    drift_client.program_id, maker_user_account.authority
+                ),
+                maker_user_account=maker_user_account,
+                order=None,
+            )
+        )
 
     # NOTE: spot markets use corresponding mint precision, not all spot market use BASE_PRECISION like perps
     spot_market_account = drift_client.get_spot_market_account(spot_market_index)
@@ -174,16 +210,17 @@ async def demo_spot_place_and_take(drift_client: DriftClient, user_map: UserMap,
                 direction=spot_trade_direction,
                 price=oracle_price.price,
             ),
-            None, # if you want to take on openbook or phoenix then need to pass in corresponding FulfillmentInfo account here
+            None,  # if you want to take on openbook or phoenix then need to pass in corresponding FulfillmentInfo account here
             maker_infos,
-            referrer_info=user_stats.get_referrer_info(), # if your UserAccount was referred, this is required.
-        )
+            referrer_info=user_stats.get_referrer_info(),  # if your UserAccount was referred, this is required.
+        ),
     ]
     tx = await drift_client.tx_sender.get_versioned_tx(
-        ixs,
-        drift_client.wallet.payer,
-        [await drift_client.fetch_market_lookup_table()])
-    tx_sim = await drift_client.connection.simulate_transaction(tx) # simulate because it's been hard to land a tx
+        ixs, drift_client.wallet.payer, [await drift_client.fetch_market_lookup_table()]
+    )
+    tx_sim = await drift_client.connection.simulate_transaction(
+        tx
+    )  # simulate because it's been hard to land a tx
     print(f"Error: {tx_sim.value.err}")
     print(f"CU used: {tx_sim.value.units_consumed}")
     print(f"logs:")
@@ -195,10 +232,9 @@ async def main():
     url = os.getenv("RPC_URL")
 
     sub_account_id = 2
-    perp_market_index = 0 # SOL-PERP
-    spot_market_index = 9 # JTO/USDC
-    trade_amount_base = 0.1 # demo buying 0.1 SOL
-
+    perp_market_index = 0  # SOL-PERP
+    spot_market_index = 9  # JTO/USDC
+    trade_amount_base = 0.1  # demo buying 0.1 SOL
 
     kp = load_keypair(secret)
     wallet = Wallet(kp)
@@ -215,19 +251,34 @@ async def main():
     await drift_client.subscribe()
 
     # user map stores User account info, this is needed to get the UserStats account for makers we want to take against
-    user_map = UserMap(UserMapConfig(
-        drift_client,
-        WebsocketConfig(),
-        drift_client.connection,
-        skip_initial_load=True, # will lazy load if True
-    ))
+    user_map = UserMap(
+        UserMapConfig(
+            drift_client,
+            WebsocketConfig(),
+            drift_client.connection,
+            skip_initial_load=True,  # will lazy load if True
+        )
+    )
     await user_map.subscribe()
 
-    await demo_perp_place_and_take(drift_client=drift_client, user_map=user_map, perp_market_index=perp_market_index, perp_trade_direction=PositionDirection.Long(), perp_base_amount=trade_amount_base)
+    await demo_perp_place_and_take(
+        drift_client=drift_client,
+        user_map=user_map,
+        perp_market_index=perp_market_index,
+        perp_trade_direction=PositionDirection.Long(),
+        perp_base_amount=trade_amount_base,
+    )
     print("")
-    await demo_spot_place_and_take(drift_client=drift_client, user_map=user_map, spot_market_index=spot_market_index, spot_trade_direction=PositionDirection.Long(), spot_base_amount=trade_amount_base)
+    await demo_spot_place_and_take(
+        drift_client=drift_client,
+        user_map=user_map,
+        spot_market_index=spot_market_index,
+        spot_trade_direction=PositionDirection.Long(),
+        spot_base_amount=trade_amount_base,
+    )
     print("")
 
+
 if __name__ == "__main__":
     asyncio.run(main())
-    print("done")
+    print("done")
diff --git a/src/driftpy/addresses.py b/src/driftpy/addresses.py
@@ -95,6 +95,7 @@ def get_prelaunch_oracle_public_key(program_id: Pubkey, market_index: int) -> Pu
         [b"prelaunch_oracle", int_to_le_bytes(market_index)], program_id
     )[0]
 
+
 def get_serum_open_orders_public_key(
     program_id: Pubkey,
     market: Pubkey,
@@ -103,15 +104,17 @@ def get_serum_open_orders_public_key(
         [b"serum_open_orders", bytes(market)], program_id
     )[0]
 
+
 def get_serum_signer_public_key(
-        program_id: Pubkey,
-        market: Pubkey,
-        nonce: int,
+    program_id: Pubkey,
+    market: Pubkey,
+    nonce: int,
 ) -> Pubkey:
     return Pubkey.create_program_address(
         [bytes(market), int_to_le_bytes(nonce)], program_id
     )
 
+
 def get_serum_fulfillment_config_public_key(
     program_id: Pubkey,
     market: Pubkey,
@@ -120,6 +123,7 @@ def get_serum_fulfillment_config_public_key(
         [b"serum_fulfillment_config", bytes(market)], program_id
     )[0]
 
+
 def get_phoenix_fulfillment_config_public_key(
     program_id: Pubkey,
     market: Pubkey,