fix: include total supply in token.categorized-supply model

credmark · Jun 18, 2024 · b2f8d6a · b2f8d6a
1 parent 7a42eb5
commit b2f8d6a
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/models/credmark/tokens/token.py b/models/credmark/tokens/token.py
@@ -762,12 +762,14 @@ class CategorizedSupplyCategory(CategorizedSupplyRequest.CategorizedSupplyCatego
         valueUsd: float = 0.0
     categories: List[CategorizedSupplyCategory]
     _iterator: str = 'categories'
+    totalSupplyScaled: float = 0.0
+    totalSupplyUsd: float = 0.0
     circulatingSupplyScaled: float = 0.0
     circulatingSupplyUsd: float = 0.0
 
 
 @Model.describe(slug='token.categorized-supply',
-                version='1.2',
+                version='1.3',
                 display_name='Token Categorized Supply',
                 description='The categorized supply for a token',
                 category='protocol',
@@ -777,7 +779,7 @@ class CategorizedSupplyCategory(CategorizedSupplyRequest.CategorizedSupplyCatego
 class TokenCirculatingSupply(Model):
     def run(self, input: CategorizedSupplyRequest) -> CategorizedSupplyResponse:
         response = CategorizedSupplyResponse(**input.dict())
-        total_supply_scaled = input.token.scaled(input.token.total_supply)
+        response.totalSupplyScaled = input.token.scaled(input.token.total_supply)
         token_price = PriceWithQuote(
             **self.context.models.price.quote({'base': input.token}))
         if token_price is None:
@@ -794,12 +796,15 @@ def run(self, input: CategorizedSupplyRequest) -> CategorizedSupplyResponse:
             categoryName='uncategorized',
             categoryType='uncategorized',
             circulating=True,
-            amountScaled=total_supply_scaled -
+            amountScaled=response.totalSupplyScaled -
             sum(c.amountScaled for c in response.categories)
         ))
         response.circulatingSupplyScaled = sum(
             c.amountScaled for c in response.categories if c.circulating)
+
         if isinstance(token_price.price, float):
+            if isinstance(response.totalSupplyScaled, float):
+                response.totalSupplyUsd = response.totalSupplyScaled * token_price.price
             if isinstance(response.circulatingSupplyScaled, float):
                 response.circulatingSupplyUsd = response.circulatingSupplyScaled * token_price.price
         return response