diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py index c7cd1259ae..10d1859fb8 100644 --- a/sdk/python/feast/client.py +++ b/sdk/python/feast/client.py @@ -40,27 +40,17 @@ from feast.online_response import OnlineResponse, _infer_online_entity_rows from feast.protos.feast.core.CoreService_pb2 import ( ApplyEntityRequest, - ApplyEntityResponse, ApplyFeatureTableRequest, - ApplyFeatureTableResponse, ArchiveProjectRequest, - ArchiveProjectResponse, CreateProjectRequest, - CreateProjectResponse, DeleteFeatureTableRequest, GetEntityRequest, - GetEntityResponse, GetFeastCoreVersionRequest, GetFeatureTableRequest, - GetFeatureTableResponse, ListEntitiesRequest, - ListEntitiesResponse, ListFeaturesRequest, - ListFeaturesResponse, ListFeatureTablesRequest, - ListFeatureTablesResponse, ListProjectsRequest, - ListProjectsResponse, ) from feast.protos.feast.core.CoreService_pb2_grpc import CoreServiceStub from feast.protos.feast.serving.ServingService_pb2 import ( @@ -388,7 +378,7 @@ def list_projects(self) -> List[str]: ListProjectsRequest(), timeout=self._config.getint(opt.GRPC_CONNECTION_TIMEOUT), metadata=self._get_grpc_metadata(), - ) # type: ListProjectsResponse + ) return list(response.projects) def create_project(self, project: str): @@ -408,7 +398,7 @@ def create_project(self, project: str): CreateProjectRequest(name=project), timeout=self._config.getint(opt.GRPC_CONNECTION_TIMEOUT), metadata=self._get_grpc_metadata(), - ) # type: CreateProjectResponse + ) def archive_project(self, project): """ @@ -430,7 +420,7 @@ def archive_project(self, project): ArchiveProjectRequest(name=project), timeout=self._config.getint(opt.GRPC_CONNECTION_TIMEOUT), metadata=self._get_grpc_metadata(), - ) # type: ArchiveProjectResponse + ) except grpc.RpcError as e: raise grpc.RpcError(e.details()) @@ -523,7 +513,7 @@ def _apply_entity(self, project: str, entity: Entity): ApplyEntityRequest(project=project, spec=entity_proto), # type: ignore timeout=self._config.getint(opt.GRPC_CONNECTION_TIMEOUT), metadata=self._get_grpc_metadata(), - ) # type: ApplyEntityResponse + ) except grpc.RpcError as e: raise grpc.RpcError(e.details()) @@ -558,7 +548,7 @@ def list_entities( # Get latest entities from Feast Core entity_protos = self._core_service.ListEntities( ListEntitiesRequest(filter=filter), metadata=self._get_grpc_metadata(), - ) # type: ListEntitiesResponse + ) # Extract entities and return entities = [] @@ -593,7 +583,7 @@ def get_entity(self, name: str, project: str = None) -> Entity: get_entity_response = self._core_service.GetEntity( GetEntityRequest(project=project, name=name.strip()), metadata=self._get_grpc_metadata(), - ) # type: GetEntityResponse + ) except grpc.RpcError as e: raise grpc.RpcError(e.details()) entity = Entity.from_proto(get_entity_response.entity) @@ -646,7 +636,7 @@ def _apply_feature_table(self, project: str, feature_table: FeatureTable): ApplyFeatureTableRequest(project=project, table_spec=feature_table_proto), # type: ignore timeout=self._config.getint(opt.GRPC_CONNECTION_TIMEOUT), metadata=self._get_grpc_metadata(), - ) # type: ApplyFeatureTableResponse + ) except grpc.RpcError as e: raise grpc.RpcError(e.details()) @@ -683,7 +673,7 @@ def list_feature_tables( feature_table_protos = self._core_service.ListFeatureTables( ListFeatureTablesRequest(filter=filter), metadata=self._get_grpc_metadata(), - ) # type: ListFeatureTablesResponse + ) # Extract feature tables and return feature_tables = [] @@ -718,7 +708,7 @@ def get_feature_table(self, name: str, project: str = None) -> FeatureTable: get_feature_table_response = self._core_service.GetFeatureTable( GetFeatureTableRequest(project=project, name=name.strip()), metadata=self._get_grpc_metadata(), - ) # type: GetFeatureTableResponse + ) except grpc.RpcError as e: raise grpc.RpcError(e.details()) return FeatureTable.from_proto(get_feature_table_response.table) @@ -785,7 +775,7 @@ def list_features_by_ref( feature_protos = self._core_service.ListFeatures( ListFeaturesRequest(filter=filter), metadata=self._get_grpc_metadata(), - ) # type: ListFeaturesResponse + ) # Extract features and return features_dict = {} diff --git a/sdk/python/feast/config.py b/sdk/python/feast/config.py index 788d257b9f..9bc8cf4913 100644 --- a/sdk/python/feast/config.py +++ b/sdk/python/feast/config.py @@ -112,8 +112,8 @@ def __init__( if options and isinstance(options, dict): self._options = options - self._config = config # type: ConfigParser - self._path = path # type: str + self._config = config + self._path = path def _get(self, option, default, get_method): fallback = {} if default is _UNSET else {"fallback": default} diff --git a/sdk/python/feast/entity.py b/sdk/python/feast/entity.py index f73e8c70ef..16740bdc17 100644 --- a/sdk/python/feast/entity.py +++ b/sdk/python/feast/entity.py @@ -49,8 +49,9 @@ def __init__( else: self._join_key = name + self._labels: MutableMapping[str, str] if labels is None: - self._labels = dict() # type: MutableMapping[str, str] + self._labels = dict() else: self._labels = labels diff --git a/sdk/python/feast/feature_table.py b/sdk/python/feast/feature_table.py index 52b97f868c..84463baa72 100644 --- a/sdk/python/feast/feature_table.py +++ b/sdk/python/feast/feature_table.py @@ -59,8 +59,10 @@ def __init__( self._features = features self._batch_source = batch_source self._stream_source = stream_source + + self._labels: MutableMapping[str, str] if labels is None: - self._labels = dict() # type: MutableMapping[str, str] + self._labels = dict() else: self._labels = labels diff --git a/sdk/python/feast/infra/provider.py b/sdk/python/feast/infra/provider.py index 2775c48173..f778032c17 100644 --- a/sdk/python/feast/infra/provider.py +++ b/sdk/python/feast/infra/provider.py @@ -172,7 +172,7 @@ def _get_requested_feature_views_to_features_dict( """Create a dict of FeatureView -> List[Feature] for all requested features. Set full_feature_names to True to have feature names prefixed by their feature view name.""" - feature_views_to_feature_map = {} # type: Dict[FeatureView, List[str]] + feature_views_to_feature_map: Dict[FeatureView, List[str]] = {} for ref in feature_refs: ref_parts = ref.split(":") diff --git a/sdk/python/tests/feast_serving_server.py b/sdk/python/tests/feast_serving_server.py index cfa2f993ef..ae7a81cd1d 100644 --- a/sdk/python/tests/feast_serving_server.py +++ b/sdk/python/tests/feast_serving_server.py @@ -5,8 +5,6 @@ import grpc -from feast.protos.feast.core import FeatureTable_pb2 as FeatureTableProto -from feast.protos.feast.core.CoreService_pb2 import ListFeatureTablesResponse from feast.protos.feast.core.CoreService_pb2_grpc import CoreServiceStub from feast.protos.feast.serving import ServingService_pb2_grpc as Serving from feast.protos.feast.serving.ServingService_pb2 import GetFeastServingInfoResponse @@ -19,9 +17,7 @@ def __init__(self, core_url: str = None): if core_url: self.__core_channel = None self.__connect_core(core_url) - self._feature_tables = ( - dict() - ) # type: Dict[str, FeatureTableProto.FeatureTable] + self._feature_tables: Dict[str, str] = (dict()) def __connect_core(self, core_url: str): if not core_url: @@ -42,9 +38,7 @@ def __connect_core(self, core_url: str): def __get_feature_tables_from_core(self): # Get updated list of feature tables - feature_tables = ( - self._core_service_stub.ListFeatureTables - ) # type: ListFeatureTablesResponse + feature_tables = self._core_service_stub.ListFeatureTables # Store each feature table locally for feature_table in list(feature_tables.tables): diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py index cf09fa236c..de92a8deb6 100644 --- a/sdk/python/tests/test_client.py +++ b/sdk/python/tests/test_client.py @@ -657,7 +657,7 @@ def test_get_online_features( entity_rows=entity_rows, feature_refs=["driver:age", "driver:rating", "driver:null_value"], project="driver_project", - ) # type: GetOnlineFeaturesResponse + ) mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with( request, metadata=auth_metadata, timeout=10 ) @@ -747,7 +747,7 @@ def test_get_online_features_multi_entities( entity_rows=entity_rows, feature_refs=["driver:age", "driver:rating", "driver:null_value"], project="driver_project", - ) # type: GetOnlineFeaturesResponse + ) mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with( request, metadata=auth_metadata, timeout=10 ) diff --git a/sdk/python/tests/test_offline_online_store_consistency.py b/sdk/python/tests/test_offline_online_store_consistency.py index 3e1cb1a9c7..fb6939af51 100644 --- a/sdk/python/tests/test_offline_online_store_consistency.py +++ b/sdk/python/tests/test_offline_online_store_consistency.py @@ -78,7 +78,7 @@ def prep_bq_fs_and_fv( df = create_dataset() job_config = bigquery.LoadJobConfig() - table_ref = f"{gcp_project}.{bigquery_dataset}.{bq_source_type}_correctness_{int(time.time())}" + table_ref = f"{gcp_project}.{bigquery_dataset}.{bq_source_type}_correctness_{int(time.time_ns())}" query = f"SELECT * FROM `{table_ref}`" job = client.load_table_from_dataframe(df, table_ref, job_config=job_config) job.result()