Fix pool_data __all__; use pool types

enigmarikki · Nov 8, 2022 · bee5302 · bee5302
1 parent 7b2c5c0
commit bee5302
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 15 deletions.
diff --git a/curvesim/pipelines/utils.py b/curvesim/pipelines/utils.py
@@ -2,6 +2,8 @@
 
 from numpy import append
 
+from curvesim.pool import MetaPool, Pool
+
 
 # Volume Multipliers
 def compute_volume_multipliers(pool_vol, market_vol, n, pool_type, mode=1):
@@ -32,10 +34,10 @@ def compute_volume_multipliers(pool_vol, market_vol, n, pool_type, mode=1):
         3: mode 2 for trades with meta-pool asset, mode 1 for basepool-only trades
 
     """
-    if pool_type == "Pool":
+    if pool_type == Pool:
         vol_mult = pool_vol_mult(pool_vol, market_vol, n, mode)
 
-    elif pool_type == "MetaPool":
+    elif pool_type == MetaPool:
         vol_mult = metapool_vol_mult(pool_vol, market_vol, n, mode)
 
     else:

diff --git a/curvesim/pool_data/__init__.py b/curvesim/pool_data/__init__.py
@@ -3,16 +3,7 @@
 Currently supports stableswap pools, meta-pools, and rebasing (RAI) metapools.
 """
 
-__all__ = [
-    "MetaPool",
-    "Pool",
-    "PoolData",
-    "RaiPool",
-    "from_address",
-    "from_symbol",
-    "get",
-    "queries",
-]
+__all__ = ["Pool", "from_address", "from_symbol", "get"]
 
 from numpy import array
 
@@ -261,17 +252,17 @@ def n(self):
 
     def type(self):
         """
-        Returns a string of the pool type.
+        Returns the pool type.
 
         Returns
         -------
         str
 
         """
         if self.dict["basepool"]:
-            return "MetaPool"
+            return MetaPool
         else:
-            return "Pool"
+            return Pool
 
     def redemption_prices(self, days=60, store=False, get_cache=True):
         """