format and lint

Signed-off-by: Achal Shah <achals@gmail.com>
feast-dev · feast-ci-bot · Jul 22, 2021 · Jul 6, 2021 · Jul 6, 2021 · Jul 6, 2021
commit 27ea33e35a1cc5be901d7da8887226880126145c
@@ -20,12 +20,12 @@
 
 import pandas as pd
 from colorama import Fore, Style
-from feast.feature_service import FeatureService
 from tqdm import tqdm
 
 from feast import utils
 from feast.entity import Entity
 from feast.errors import FeastProviderLoginError, FeatureViewNotFoundException
+from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
 from feast.inference import (
     update_data_sources_with_inferred_event_timestamp_col,
@@ -256,7 +256,9 @@ def apply(
 
     @log_exceptions_and_usage
     def get_historical_features(
-        self, entity_df: Union[pd.DataFrame, str], feature_refs: Union[List[str], List[FeatureService]],
+        self,
+        entity_df: Union[pd.DataFrame, str],
+        feature_refs: Union[List[str], List[FeatureService]],
     ) -> RetrievalJob:
         """Enrich an entity dataframe with historical feature values for either training or batch scoring.
 
@@ -297,7 +299,9 @@ def get_historical_features(
         """
 
         if not len(feature_refs) > 0:
-            raise ValueError("A list of feature references or FeatureService must be provided.")
+            raise ValueError(
+                "A list of feature references or FeatureService must be provided."
+            )
 
         all_feature_views = self._registry.list_feature_views(project=self.project)
         try:
@@ -307,15 +311,18 @@ def get_historical_features(
         except FeatureViewNotFoundException as e:
             sys.exit(e)
 
+        _feature_refs: List[str]
         if isinstance(feature_refs[0], FeatureService):
-            feature_refs = _get_features_refs_from_feature_services(feature_refs)
+            _feature_refs = _get_features_refs_from_feature_services(feature_refs)  # type: ignore
+        else:
+            _feature_refs = feature_refs  # type: ignore
 
         provider = self._get_provider()
         try:
             job = provider.get_historical_features(
                 self.config,
                 feature_views,
-                feature_refs,
+                _feature_refs,
                 entity_df,
                 self._registry,
                 self.project,
@@ -487,7 +494,9 @@ def tqdm_builder(length):
 
     @log_exceptions_and_usage
     def get_online_features(
-        self, feature_refs: Union[List[str], List[FeatureService]], entity_rows: List[Dict[str, Any]],
+        self,
+        feature_refs: Union[List[str], List[FeatureService]],
+        entity_rows: List[Dict[str, Any]],
     ) -> OnlineResponse:
         """
         Retrieves the latest online feature data.
@@ -607,7 +616,8 @@ def _entity_row_to_field_values(
 
 
 def _group_refs(
-    feature_refs: Union[List[str], FeatureService], all_feature_views: List[FeatureView]
+    feature_refs: Union[List[str], List[FeatureService]],
+    all_feature_views: List[FeatureView],
 ) -> List[Tuple[FeatureView, List[str]]]:
     """ Get list of feature views and corresponding feature names based on feature references"""
 
@@ -617,23 +627,29 @@ def _group_refs(
     # view name to feature names
     views_features = defaultdict(list)
 
-    if isinstance(feature_refs, list):
+    if isinstance(feature_refs[0], str):
         for ref in feature_refs:
-            view_name, feat_name = ref.split(":")
+            view_name, feat_name = ref.split(":")  # type: ignore
             if view_name not in view_index:
                 raise FeatureViewNotFoundException(view_name)
             views_features[view_name].append(feat_name)
-    if isinstance(feature_refs, FeatureService):
-        for feature_projection in feature_refs.features:
-            views_features[feature_projection.name].extend([f.name for f in feature_projection.features])
+    elif isinstance(feature_refs[0], FeatureService):
+        feature_services: List[FeatureService] = feature_refs  # type: ignore
+        for feature_service in feature_services:
+            for feature_projection in feature_service.features:
+                views_features[feature_projection.name].extend(
+                    [f.name for f in feature_projection.features]
+                )
 
     result = []
     for view_name, feature_names in views_features.items():
         result.append((view_index[view_name], feature_names))
     return result
 
 
-def _get_features_refs_from_feature_services(feature_services: List[FeatureService]) -> List[str]:
+def _get_features_refs_from_feature_services(
+    feature_services: List[FeatureService],
+) -> List[str]:
     feature_refs = []
     for fs in feature_services:
         name = fs.name
@@ -642,7 +658,8 @@ def _get_features_refs_from_feature_services(feature_services: List[FeatureServi
 
 
 def _get_requested_feature_views(
-    feature_refs: Union[List[str], FeatureService], all_feature_views: List[FeatureView]
+    feature_refs: Union[List[str], List[FeatureService]],
+    all_feature_views: List[FeatureView],
 ) -> List[FeatureView]:
     """Get list of feature views based on feature references"""
     # TODO: Get rid of this function. We only need _group_refs