2 files changed, 21 insertions, 9 deletions
diff --git a/python/book_bbg.py b/python/book_bbg.py
index 54212412..e12646cb 100644
--- a/python/book_bbg.py
+++ b/python/book_bbg.py
@@ -1,15 +1,11 @@
 from serenitas.utils.remote import SftpClient
-import datetime
+from trade_dataclasses import Deal, DealType, BbgDeal
 import csv
-from trade_dataclasses import CDSDeal, BondDeal, BbgDeal, _funds, _cdx_cp, _fcms
-from decimal import Decimal
 from stat import S_ISREG
 
-_deal = {"BOND": BondDeal, "CDX": CDSDeal}
-
 
 def trade_booking_process(file_handle, index, asset_class):
-    deal = _deal[asset_class]
+    deal = Deal[DealType(asset_class)]
     for row in csv.DictReader(file_handle):
         line = {"bbg_ticket_id": index, **row}
         trade = deal.from_bbg_line(line)
diff --git a/python/trade_dataclasses.py b/python/trade_dataclasses.py
index 08f92638..e3790466 100644
--- a/python/trade_dataclasses.py
+++ b/python/trade_dataclasses.py
@@ -149,16 +149,26 @@ class Counterparty:
 register_adapter(Frequency, lambda f: AsIs(f.value))
 
 
+class DealType(Enum):
+    Bond = "BOND"
+    CDS = "CDX"
+
+
 class Deal:
     _conn: ClassVar = dbconn("dawndb", application_name="autobooker")
+    _registry = {}
     _table_name: None
     _sql_fields: ClassVar[list[str]]
     _sql_insert: ClassVar[str]
     _sql_select: ClassVar[str]
     _insert_queue: ClassVar[list] = []
 
-    def __init_subclass__(cls, table_name: str, insert_ignore=()):
+    def __class_getitem__(cls, deal_type: DealType):
+        return cls._registry[deal_type]
+
+    def __init_subclass__(cls, deal_type: DealType, table_name: str, insert_ignore=()):
         super().__init_subclass__()
+        cls._registry[deal_type] = cls
         cls._table_name = table_name
         insert_columns = [c for c in cls.__annotations__ if c not in insert_ignore]
         place_holders = ",".join(["%s"] * len(insert_columns))
@@ -230,7 +240,13 @@ class BbgDeal:
 
 
 @dataclass
-class CDSDeal(BbgDeal, Deal, table_name="cds", insert_ignore=("id", "dealid")):
+class CDSDeal(
+    BbgDeal,
+    Deal,
+    deal_type=DealType.CDS,
+    table_name="cds",
+    insert_ignore=("id", "dealid"),
+):
     fund: Fund = field(metadata={"mtm": "Account Abbreviation"})
     account_code: str
     cp_code: str = field(metadata={"mtm": "Broker Id"})
@@ -321,7 +337,7 @@ class CDSDeal(BbgDeal, Deal, table_name="cds", insert_ignore=("id", "dealid")):
 
 
 @dataclass
-class BondDeal(BbgDeal, Deal, table_name="bonds"):
+class BondDeal(BbgDeal, Deal, deal_type=DealType.Bond, table_name="bonds"):
     buysell: bool
     description: str
     faceamount: float