market_trade/core/dealManager.py

@@ -1,49 +1,77 @@
 import pandas as pd
 import datetime
 import numpy as np
-import uuid 
+import uuid
+
 
 class DealManager():
-    
+    """Manages open trading positions and deal lifecycle.
+
+    Tracks active positions with their entry prices and quantities,
+    supporting both opening new positions and closing existing ones.
+    """
+
     def __init__(self):
-        #self.commission=0.04
-        self.columns=['uuid','figi','amount','startPrice']
+        """Initialize DealManager with empty deals DataFrame."""
+        self.columns = ['uuid', 'figi', 'amount', 'startPrice']
         self.deals = pd.DataFrame(columns=self.columns)
         self.deals = self.deals.set_index('uuid')
-        
-    def findDealByPriceAndFig(self,price,figi):
-        ans=None
+
+    def find_deal_by_price_and_figi(self, price: float, figi: str):
+        """Find existing deal by price and instrument identifier.
+
+        Args:
+            price: Entry price to search for.
+            figi: Financial Instrument Global Identifier.
+
+        Returns:
+            Deal UUID if found, None otherwise.
+        """
+        ans = None
         for i in range(self.deals.shape[0]):
             if self.deals.iloc[i].startPrice == price and self.deals.iloc[i].figi == figi:
                 ans = self.deals.iloc[i].name
                 break
         return ans
 
-    def openDeal(self,figi,startPrice,amount=1):
-        desiredDeal=self.findDealByPriceAndFig(startPrice,figi)
-        if desiredDeal == None:
-            newDealDict={
-                'uuid':[str(uuid.uuid4())],
-                'figi':[figi],
-                'startPrice':[startPrice],
-                'amount':[amount]
+    def open_deal(self, figi: str, start_price: float, amount: int = 1) -> None:
+        """Open new deal or add to existing position.
+
+        If a deal with the same FIGI and price exists, adds to the amount.
+        Otherwise creates a new deal entry.
+
+        Args:
+            figi: Financial Instrument Global Identifier.
+            start_price: Entry price for the position.
+            amount: Number of units to trade (default 1).
+        """
+        desired_deal = self.find_deal_by_price_and_figi(start_price, figi)
+        if desired_deal is None:
+            new_deal_dict = {
+                'uuid': [str(uuid.uuid4())],
+                'figi': [figi],
+                'startPrice': [start_price],
+                'amount': [amount]
             }
 
-            #newDealDict['profit']=[startPrice*pow(1+self.commission,2)]
-
-
-
-            newDeal=pd.DataFrame.from_dict(newDealDict).set_index('uuid')
-            self.deals=pd.concat([self.deals, newDeal])
+            new_deal = pd.DataFrame.from_dict(new_deal_dict).set_index('uuid')
+            self.deals = pd.concat([self.deals, new_deal])
         else:
-            self.deals.at[desiredDeal,'amount'] += amount
+            self.deals.at[desired_deal, 'amount'] += amount
 
-    def closeDeal(self,uuid,amount):
-        
-        desiredDeal=self.deals.loc[uuid]
-        if desiredDeal.amount - amount == 0:
-            self.deals = self.deals.drop(labels = [uuid],axis = 0)
+    def close_deal(self, uuid_str: str, amount: int) -> None:
+        """Close deal partially or completely.
+
+        Args:
+            uuid_str: Deal UUID to close.
+            amount: Number of units to close.
+
+        Note:
+            If amount equals deal amount, removes deal entirely.
+            Otherwise decreases deal amount.
+        """
+        desired_deal = self.deals.loc[uuid_str]
+        if desired_deal.amount - amount == 0:
+            self.deals = self.deals.drop(labels=[uuid_str], axis=0)
         else:
-            self.deals.at[uuid,'amount'] -= amount
-            #self.deals.loc[uuid].amount = desiredDeal.amount - amount
-        
+            self.deals.at[uuid_str, 'amount'] -= amount