Merge pull request #141 from AsyncAlgoTrading/instattr

Better instrument handling of multiple exchanges

aat/core/instrument/calendar.py

@@ -70,3 +70,8 @@ def open(self) -> Optional[Tuple[time, ...]]:
     @property
     def close(self) -> Optional[Tuple[time, ...]]:
         return self._close_times
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TradingDay):
+            return False
+        return self.open == other.open and self.close == other.close

aat/core/instrument/db.py

@@ -11,14 +11,25 @@ class InstrumentDB(object):
     """instrument registration"""
 
     def __init__(self) -> None:
-        self._name_map: Dict[str, "Instrument"] = {}
-        self._map: Dict[Tuple[str, InstrumentType], "Instrument"] = {}
+        self._by_name: Dict[str, List["Instrument"]] = {}
+        self._by_type: Dict[Tuple[str, InstrumentType], List["Instrument"]] = {}
+        self._by_exchange: Dict[Tuple[str, ExchangeType], "Instrument"] = {}
+        self._by_type_and_exchange: Dict[
+            Tuple[str, InstrumentType, ExchangeType], "Instrument"
+        ] = {}
 
     def add(self, instrument: "Instrument") -> None:
-        if instrument.name in self._name_map:
-            return
-        self._name_map[instrument.name] = instrument
-        self._map[instrument.name, instrument.type] = instrument
+        if instrument.name not in self._by_name:
+            self._by_name[instrument.name] = [instrument]
+            self._by_type[instrument.name, instrument.type] = [instrument]
+        else:
+            self._by_name[instrument.name].append(instrument)
+            self._by_type[instrument.name, instrument.type].append(instrument)
+
+        self._by_exchange[instrument.name, instrument.exchange] = instrument
+        self._by_type_and_exchange[
+            instrument.name, instrument.type, instrument.exchange
+        ] = instrument
 
     def instruments(
         self,
@@ -28,13 +39,11 @@ def instruments(
         *args: Tuple,
         **kwargs: Dict
     ) -> List["Instrument"]:
+        ret = [inst for values in self._by_name.values() for inst in values]
         if not name and not type and not exchange:
-            return list(self._map.values())
-        elif name:
-            inst = self._name_map.get(name, None)
-            return [inst] if inst else []
-
-        ret = list(self._map.values())
+            return ret
+        if name:
+            ret = [r for r in ret if r.name == name]
         if type:
             ret = [r for r in ret if r.type == type]
         if exchange:
@@ -45,11 +54,16 @@ def get(
         self,
         name: str = "",
         type: InstrumentType = InstrumentType.EQUITY,
-        exchange: ExchangeType = ExchangeType(""),
+        exchange: Optional[ExchangeType] = ExchangeType(""),
         *args: Tuple,
         **kwargs: Dict
-    ) -> "Instrument":
-        return self._name_map[name]
-
-
-# TODO allow for multiple exchange's distinct instrument representation
+    ) -> Optional["Instrument"]:
+        """Like `instruments` but only returns 1"""
+        ret = [inst for values in self._by_name.values() for inst in values]
+        if name:
+            ret = [r for r in ret if r.name == name]
+        if type:
+            ret = [r for r in ret if r.type == type]
+        if exchange:
+            ret = [r for r in ret if exchange in r.exchanges]
+        return ret[0] if ret else None

aat/core/instrument/instrument.py

@@ -12,9 +12,10 @@
 class Instrument(object):
     _instrumentdb = InstrumentDB()
 
+    __exchange: ExchangeType
     __exchanges: List[ExchangeType]
     __type: InstrumentType
-    __trading_days: Dict[ExchangeType, TradingDay]
+    __trading_day: TradingDay
     __broker_exchange: Optional[ExchangeType]
     __broker_id: Optional[str]
     __currency: Optional["Instrument"]
@@ -32,8 +33,9 @@ class Instrument(object):
     __slots__ = [
         "__name",
         "__type",
+        "__exchange",
         "__exchanges",
-        "__trading_days",
+        "__trading_day",
         "__broker_exchange",
         "__broker_id",
         "__currency",
@@ -49,7 +51,7 @@ class Instrument(object):
     ]
 
     def __new__(cls: Type, *args: Tuple, **kwargs: Dict) -> "Instrument":
-        if cls._instrumentdb.instruments(*args, **kwargs):
+        if cls._instrumentdb.get(*args, **kwargs):
             return cls._instrumentdb.get(*args, **kwargs)
 
         if _CPP:
@@ -140,43 +142,37 @@ def __init__(
         ):
             # append this exchange to list of available
             self.__exchanges.append(exchange)
+
+            # set attribute
+            self.__exchange = exchange
+
         elif exchange:
             # create new list with this one
             self.__exchanges = [exchange]
+
+            # set attribute
+            self.__exchange = exchange
+
         elif hasattr(self, "_Instrument__exchanges"):
             # do nothing
             pass
+
         else:
             # no exchange known and no exchange provided
             self.__exchanges = []
 
+            # set attribute
+            self.__exchange = exchange
+
         # Optional Fields
-        if (
-            exchange
-            and kwargs.get("trading_day", None)
-            and hasattr(self, "_Instrument__trading_days")
-            and exchange not in self.__trading_days
-        ):
-            # append this exchange to list of available
-            self.__trading_days[exchange] = cast(TradingDay, kwargs.get("trading_day"))
-        elif (
-            exchange
-            and kwargs.get("trading_day", None)
-            and not hasattr(self, "_Instrument__trading_days")
-        ):
-            # create new dict with this one
-            self.__trading_days = {
-                exchange: cast(TradingDay, kwargs.get("trading_day"))
-            }
-        elif kwargs.get("trading_day", None):
-            # no exchange, raise error
-            raise Exception("If setting trading hours, must provide exchange")
-        elif hasattr(self, "_Instrument__trading_days"):
-            # do nothing
-            pass
+        if hasattr(self, "_Instrument__trading_day") and self.__trading_day is not None:
+            assert kwargs.get(
+                "trading_day"
+            ) is None or self.__trading_day == kwargs.get("trading_day")
         else:
-            # no exchange known and no trasding days provided
-            self.__trading_days = {}
+            self.__trading_day: Optional[TradingDay] = cast(
+                TradingDay, kwargs.get("trading_day")
+            )
 
         if (
             hasattr(self, "_Instrument__broker_exchange")
@@ -303,15 +299,35 @@ def type(self) -> InstrumentType:
     def exchanges(self) -> List[ExchangeType]:
         return self.__exchanges
 
+    @property
+    def exchange(self) -> ExchangeType:
+        return self.__exchange
+
+    def tradingLines(self, exchange: ExchangeType = None) -> List["Instrument"]:
+        """Returns other exchanges that the same instrument trades on
+
+        Returns:
+            exchange (ExchangeType): Exchange to filter by
+        """
+        return self._instrumentdb.instruments(self.name, self.type, exchange)
+
+    def synthetics(
+        self, type: InstrumentType = None, exchange: ExchangeType = None
+    ) -> List["Instrument"]:
+        """Returns other instruments with the same name
+
+        Returns:
+            type (InstrumentType): instrument type to filter by, e.g. for an equity, filter to only get ADRs
+            exchange (ExchangeType): Exchange to filter by
+        """
+        return self._instrumentdb.instruments(self.name, type, exchange)
+
     # ******** #
     # Optional #
     # ******** #
     @property
-    def tradingDays(self) -> Dict[ExchangeType, TradingDay]:
-        return self.__trading_days
-
-    def tradingDay(self, exchange: ExchangeType) -> TradingDay:
-        return self.__trading_days[exchange]
+    def tradingDay(self) -> TradingDay:
+        return self.__trading_day
 
     @property
     def brokerExchange(self) -> Optional[ExchangeType]:

aat/tests/core/instrument/test_calendar.py

@@ -12,9 +12,9 @@ def test_instrument_calendar_getter(self):
         e = ExchangeType("test-exchange")
 
         i = Instrument(
-            "test",
+            "TestTradingDayInst",
             exchange=e,
             trading_day=t,
         )
 
-        assert i.tradingDay(e) == t
+        assert i.tradingDay == t

aat/tests/core/instrument/test_instrument.py

@@ -0,0 +1,33 @@
+import pytest
+from aat.core import ExchangeType, Instrument
+
+
+class TestInstrument(object):
+    @pytest.mark.skipif("os.environ.get('AAT_USE_CPP', '')")
+    def test_instrument(self):
+        E1 = ExchangeType("E1")
+        E2 = ExchangeType("E2")
+        E3 = ExchangeType("E3")
+
+        i1 = Instrument(
+            "TestInst",
+            exchange=E1,
+            broker_id="1",
+        )
+
+        i2 = Instrument(
+            "TestInst",
+            exchange=E2,
+            broker_id="2",
+            broker_exchange="test",
+        )
+
+        i3 = Instrument(
+            "TestInst",
+            exchange=E3,
+            broker_id="3",
+        )
+
+        assert i1.tradingLines() == [i1, i2, i3]
+        assert i2.tradingLines() == [i1, i2, i3]
+        assert i3.tradingLines() == [i1, i2, i3]