Skip to content

Commit

Permalink
add some missing user methods such as get_perp_position
Browse files Browse the repository at this point in the history
  • Loading branch information
crispheaney committed Nov 28, 2023
1 parent 8a70fab commit 9ab81da
Showing 1 changed file with 41 additions and 3 deletions.
44 changes: 41 additions & 3 deletions src/driftpy/drift_user.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,32 @@ def get_spot_market_account(self, market_index: int) -> SpotMarketAccount:
def get_user_account(self) -> UserAccount:
return self.account_subscriber.get_user_account_and_slot().data

def get_token_amount(self, market_index: int) -> int:
spot_position = self.get_spot_position(market_index)
if spot_position is None:
return 0

spot_market = self.get_spot_market_account(market_index)
return get_token_amount(
spot_position.scaled_balance, spot_market, spot_position.balance_type
)

def get_order(self, order_id: int) -> Optional[Order]:
for order in self.get_user_account().orders:
if order.order_id == order_id:
return order

return None

def get_order_by_user_order_id(self, user_order_id: int):
for order in self.get_user_account().orders:
if order.user_order_id == user_order_id:
return order

return None

def get_open_orders(
self,
# market_type: MarketType,
# market_index: int,
# position_direction: PositionDirection
):
return list(
filter(
Expand All @@ -83,6 +104,23 @@ def get_open_orders(
)
)

def get_perp_position(self, market_index: int) -> Optional[PerpPosition]:
for position in self.get_user_account().perp_positions:
if position.market_index == market_index and not is_available(position):
return position

return None

def get_spot_position(self, market_index: int) -> Optional[SpotPosition]:
for position in self.get_user_account().spot_positions:
if (
position.market_index == market_index
and not is_spot_position_available(position)
):
return position

return None

def get_spot_market_liability(
self,
market_index=None,
Expand Down

0 comments on commit 9ab81da

Please sign in to comment.