1 files changed, 19 insertions, 3 deletions
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