Expand functionality in ExportApi.signal_query (#91)

Author: grotmol

exabel_data_sdk/client/api/export_api.py

@@ -84,10 +84,13 @@ def signal_query(
         self,
         signal: Union[str, Column, Sequence[Union[str, Column]]],
         bloomberg_ticker: Union[str, Sequence[str]] = None,
+        *,
         factset_id: Union[str, Sequence[str]] = None,
         tag: Union[str, Sequence[str]] = None,
         start_time: Union[str, pd.Timestamp] = None,
         end_time: Union[str, pd.Timestamp] = None,
+        identifier: Union[Column, Sequence[Column]] = None,
+        version: Union[str, pd.Timestamp, Sequence[str], Sequence[pd.Timestamp]] = None,
     ) -> Union[pd.Series, pd.DataFrame]:
         """
         Run a query for one or more signals.
@@ -106,6 +109,15 @@ def signal_query(
                         or with any of the provided tags if several.
             start_time: the first date to retrieve data for
             end_time:   the last date to retrieve data for
+            identifier: the identifier(s) to return to identify the entities in the result.
+                        By default, will use Signals.NAME if multiple identifiers
+                        or a tag are given, or else no identifier.
+            version:    the Point-in-Time at which to evaluate the signals,
+                        or a list of such dates at which to evaluate the signals.
+                        If a list of dates is provided, 'version' will be included as a column
+                        in the result.
+                        If no version is specified, then the signals will be evaluated with the
+                        latest available data (the default).
 
         Returns:
             A pandas Series if the result is a single time series,
@@ -116,6 +128,8 @@ def signal_query(
         """
         if not signal:
             raise ValueError("Must specify signal to retrieve")
+
+        # Specify entity filter
         multi_company = False
         predicates: List[Predicate] = []
         if factset_id:
@@ -138,18 +152,41 @@ def signal_query(
             multi_company = True
         if len(predicates) > 1:
             raise ValueError("At most one company identification method can be specified")
-        if multi_company:
-            index = [Signals.NAME, Signals.TIME]
+
+        # Specify the identifier(s)
+        index: List[Column] = []
+        if identifier is None:
+            if multi_company:
+                index.append(Signals.NAME)
+        elif isinstance(identifier, Column):
+            index.append(identifier)
         else:
-            index = [Signals.TIME]
+            index.extend(identifier)
+        index.append(Signals.TIME)
+
+        # Specify version(s), if provided
+        if version:
+            if isinstance(version, (str, pd.Timestamp)):
+                predicates.append(Signals.VERSION.equal(version))
+            else:
+                predicates.append(Signals.VERSION.in_list(*version))
+                index.insert(0, Signals.VERSION)
+
+        # The columns to query for
         columns: List[Union[str, Column]] = list(index)
         if isinstance(signal, (str, Column)):
             columns.append(signal)
         else:
             columns.extend(signal)
+
+        # Execute the query
         query = Signals.query(
             columns, start_time=start_time, end_time=end_time, predicates=predicates
         )
         df = self.run_query(query.sql())
+
+        # Set the row index
         df.set_index([col.name for col in index], inplace=True)
+        # Squeeze to a Series if a single time series was returned,
+        # and fix the data type (the backend returns a DataFrame with dtype=object)
         return df.squeeze(axis=1).infer_objects()

exabel_data_sdk/query/signals.py

@@ -16,11 +16,19 @@ class Signals:
 
     TABLE = Table("signals")
 
+    # Special columns. Refer to online documentation for descriptions:
+    # https://help.exabel.com/docs/exporting-via-exabel-sdk
     TIME = Column("time")
-    EXABEL_ID = Column("exabel_id")
-    FACTSET_ID = Column("factset_id")
-    BLOOMBERG_TICKER = Column("bloomberg_ticker")
+    VERSION = Column("version")
+    LABEL = Column("label")
     NAME = Column("name")
+    BLOOMBERG_TICKER = Column("bloomberg_ticker")
+    FACTSET_ID = Column("factset_id")
+    ISIN = Column("isin")
+    MIC = Column("mic")
+    TICKER = Column("ticker")
+    EXABEL_ID = Column("exabel_id")
+    RESOURCE_NAME = Column("resource_name")
 
     @staticmethod
     def query(

exabel_data_sdk/tests/client/api/test_export_api.py

@@ -1,8 +1,11 @@
 import unittest
 from unittest.mock import MagicMock
 
+import pandas as pd
+
 from exabel_data_sdk.client.api.export_api import ExportApi
 from exabel_data_sdk.query.column import Column
+from exabel_data_sdk.query.signals import Signals
 
 
 class TestExportApi(unittest.TestCase):
@@ -19,12 +22,20 @@ def test_signal_query(self):
             "SELECT time, Sales_Actual FROM signals WHERE factset_id = 'QLGSL2-R'"
         )
 
-        export_api.signal_query("Sales_Actual", bloomberg_ticker=["AAPL US", "TSLA US"])
+        export_api.signal_query("Sales_Actual", ["AAPL US", "TSLA US"])
         mock.assert_called_with(
             "SELECT name, time, Sales_Actual FROM signals "
             "WHERE bloomberg_ticker IN ('AAPL US', 'TSLA US')"
         )
 
+        export_api.signal_query(
+            "Sales_Actual", ["AAPL US", "TSLA US"], identifier=Signals.EXABEL_ID
+        )
+        mock.assert_called_with(
+            "SELECT exabel_id, time, Sales_Actual FROM signals "
+            "WHERE bloomberg_ticker IN ('AAPL US', 'TSLA US')"
+        )
+
         export_api.signal_query("Sales_Actual", tag="graph:tag:user:xyz-123")
         mock.assert_called_with(
             "SELECT name, time, Sales_Actual FROM signals WHERE has_tag('graph:tag:user:xyz-123')"
@@ -36,3 +47,26 @@ def test_signal_query(self):
             "SELECT time, Sales_Actual_fiscal, 'sales_actual(alignment=''afp'')' AS Q "
             "FROM signals WHERE bloomberg_ticker = 'AAPL US'"
         )
+
+        export_api.signal_query(
+            Column("Prediction", "allocations(analysis=7)"), identifier=Signals.NAME
+        )
+        mock.assert_called_with(
+            "SELECT name, time, 'allocations(analysis=7)' AS Prediction FROM signals"
+        )
+
+        export_api.signal_query("Sales_Actual", factset_id="QLGSL2-R", version="2019-02-03")
+        mock.assert_called_with(
+            "SELECT time, Sales_Actual FROM signals WHERE factset_id = 'QLGSL2-R'"
+            " AND version = '2019-02-03'"
+        )
+
+        export_api.signal_query(
+            "Sales_Actual",
+            factset_id=["QLGSL2-R", "FOOBAR"],
+            version=[pd.Timestamp("2019-01-01"), pd.Timestamp("2019-02-02")],
+        )
+        mock.assert_called_with(
+            "SELECT version, name, time, Sales_Actual FROM signals "
+            "WHERE factset_id IN ('QLGSL2-R', 'FOOBAR') AND version IN ('2019-01-01', '2019-02-02')"
+        )