specify type for dataframe (minor change)

shlomiku · Mar 21, 2021 · 5add7a5 · 5add7a5
1 parent 3d7eacd
commit 5add7a5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/zipline/data/psql_daily_bars.py b/zipline/data/psql_daily_bars.py
@@ -711,7 +711,7 @@ def _ensure_sessions_consistency(self, data_slice, invalid_data_behavior):
         return val
 
     @expect_element(invalid_data_behavior={'warn', 'raise', 'ignore'})
-    def _write_to_postgres(self, sid, data, invalid_data_behavior):
+    def _write_to_postgres(self, sid, data: pd.DataFrame, invalid_data_behavior):
 
         result = self._format_df_columns_and_index(data, sid)
         if not result.empty:
@@ -790,7 +790,7 @@ def _get_exisiting_data_dates_from_db(self, sid):
         )
         return edge_days
 
-    def _format_df_columns_and_index(self, data, sid):
+    def _format_df_columns_and_index(self, data: pd.DataFrame, sid):
         """
         make sure that the data received is in the structure we expect columns and index wise.
         :param data: data from data bundle