1 files changed, 15 insertions, 10 deletions
diff --git a/python/collateral/citi.py b/python/collateral/citi.py
index e9acfa88..7433547a 100644
--- a/python/collateral/citi.py
+++ b/python/collateral/citi.py
@@ -1,12 +1,13 @@
+import datetime
 import pandas as pd
 from . import DAILY_DIR
 from .common import load_pdf, get_col, next_business_day, parse_num
 
 
-def load_file(d):
+def load_file(d, fund="Serenitas"):
     try:
         fname = next(
-            (DAILY_DIR / "CITI_reports").glob(
+            (DAILY_DIR / fund / "CITI_reports").glob(
                 f"262966_Portfolio_{d.strftime('%Y%m%d')}*"
             )
         )
@@ -15,11 +16,13 @@ def load_file(d):
     return pd.read_excel(fname, skiprows=6, skipfooter=2)
 
 
-def download_files(em, count=20, **kwargs):
+def download_files(em, count=20, fund="Serenitas", **kwargs):
+    if fund != "Serenitas":
+        return
     emails = em.get_msgs(
         path=["NYops", "Margin Calls Citi"], count=count, subject__startswith="262966"
     )
-    DATA_DIR = DAILY_DIR / "CITI_reports"
+    DATA_DIR = DAILY_DIR / fund / "CITI_reports"
     for msg in emails:
         for attach in msg.attachments:
             fname = attach.name
@@ -38,15 +41,15 @@ def get_df(l, col1, col2, col3):
     return df
 
 
-def get_total_collateral(d):
+def get_total_collateral(d: datetime.date, fund="Serenitas"):
     try:
         fname = next(
-            (DAILY_DIR / "CITI_reports").glob(
+            (DAILY_DIR / fund / "CITI_reports").glob(
                 f"262966_MarginNotice_{d.strftime('%Y%m%d')}_*.pdf"
             )
         )
     except StopIteration:
-        raise FileNotFoundError(f"CITI file not found for date {d.date()}")
+        raise FileNotFoundError(f"CITI file not found for date {d}")
     l = load_pdf(fname)
     col1 = (370, 500, 70, 250)
     col2 = (370, 500, 300, 530)
@@ -68,9 +71,11 @@ def get_total_collateral(d):
     )
 
 
-def collateral(d, dawn_trades, **kwargs):
-    df = load_file(next_business_day(d))
-    collat = sum(get_total_collateral(d)[:2])
+def collateral(d, dawn_trades, fund="Serenitas", **kwargs):
+    if fund != "Serenitas":
+        raise ValueError
+    df = load_file(next_business_day(d), fund)
+    collat = sum(get_total_collateral(d, fund)[:2])
     df = df[["Operations File", "Market Value", "BasicAmt"]].dropna(
         subset=["Operations File"]
     )  # missing Operations File means assignment usually