CR updates

Signed-off-by: Achal Shah <achals@gmail.com>

sdk/python/feast/diff/FcoDiff.py

@@ -1,17 +1,13 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, List, Set, Tuple, Union, TypeVar, Iterable
-from unittest.mock import Base
+from typing import Any, Iterable, List, Set, Tuple, TypeVar
 
 from feast.base_feature_view import BaseFeatureView
 from feast.entity import Entity
 from feast.feature_service import FeatureService
 from feast.feature_table import FeatureTable
-from feast.feature_view import FeatureView
-from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
 from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
-from feast.request_feature_view import RequestFeatureView
 
 
 @dataclass
@@ -48,41 +44,33 @@ def add_fco_diff(self, fco_diff: FcoDiff):
         self.fco_diffs.append(fco_diff)
 
 
-T = TypeVar('T', Entity, BaseFeatureView, FeatureService, FeatureTable)
+T = TypeVar("T", Entity, BaseFeatureView, FeatureService, FeatureTable)
 
 
 def tag_objects_for_keep_delete_add(
     existing_objs: Iterable[T], desired_objs: Iterable[T]
-) -> Tuple[Iterable[T], Iterable[T], Iterable[T]]:
+) -> Tuple[Set[T], Set[T], Set[T]]:
     existing_obj_names = {e.name for e in existing_objs}
     desired_obj_names = {e.name for e in desired_objs}
 
-    objs_to_add = {
-        e for e in desired_objs if e.name not in existing_obj_names
-    }
+    objs_to_add = {e for e in desired_objs if e.name not in existing_obj_names}
     objs_to_keep = {e for e in desired_objs if e.name in existing_obj_names}
-    objs_to_delete = {
-        e for e in existing_objs if e.name not in desired_obj_names
-    }
+    objs_to_delete = {e for e in existing_objs if e.name not in desired_obj_names}
 
     return objs_to_keep, objs_to_delete, objs_to_add
 
 
-U = TypeVar('U', EntityProto, FeatureViewProto)
+U = TypeVar("U", EntityProto, FeatureViewProto)
 
 
 def tag_proto_objects_for_keep_delete_add(
-        existing_objs: Iterable[U], desired_objs: Iterable[U]
+    existing_objs: Iterable[U], desired_objs: Iterable[U]
 ) -> Tuple[Iterable[U], Iterable[U], Iterable[U]]:
     existing_obj_names = {e.spec.name for e in existing_objs}
     desired_obj_names = {e.spec.name for e in desired_objs}
 
-    objs_to_add = [
-        e for e in desired_objs if e.spec.name not in existing_obj_names
-    ]
+    objs_to_add = [e for e in desired_objs if e.spec.name not in existing_obj_names]
     objs_to_keep = [e for e in desired_objs if e.spec.name in existing_obj_names]
-    objs_to_delete = [
-        e for e in existing_objs if e.spec.name not in desired_obj_names
-    ]
+    objs_to_delete = [e for e in existing_objs if e.spec.name not in desired_obj_names]
 
     return objs_to_keep, objs_to_delete, objs_to_add

sdk/python/feast/registry.py

@@ -143,8 +143,7 @@ def diff_between(
             entities_to_delete,
             entities_to_add,
         ) = tag_proto_objects_for_keep_delete_add(
-            current_registry.entities,
-            new_registry.entities,
+            current_registry.entities, new_registry.entities,
         )
 
         for e in entities_to_add:

sdk/python/feast/repo_operations.py

@@ -161,7 +161,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
         set(registry.list_entities(project=project)), repo.entities
     )
     # TODO(achals): This code path should be refactored to handle added & kept entities separately.
-    entities_to_keep = entities_to_keep.union(entities_to_add)
+    entities_to_keep = set(entities_to_keep).union(entities_to_add)
 
     views = tag_objects_for_keep_delete_add(
         set(registry.list_feature_views(project=project)), repo.feature_views