From 5e6429f4d98e0a60c46faa27a3e95767d87364dd Mon Sep 17 00:00:00 2001 From: tokoko Date: Wed, 27 Mar 2024 00:24:13 +0000 Subject: [PATCH] remove proto-plus imports Signed-off-by: tokoko --- sdk/python/feast/base_feature_view.py | 15 ++++++++++++--- sdk/python/feast/infra/registry/base_registry.py | 2 +- sdk/python/feast/infra/registry/registry.py | 2 +- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/base_feature_view.py b/sdk/python/feast/base_feature_view.py index 975537a394..31140e2899 100644 --- a/sdk/python/feast/base_feature_view.py +++ b/sdk/python/feast/base_feature_view.py @@ -13,13 +13,20 @@ # limitations under the License. from abc import ABC, abstractmethod from datetime import datetime -from typing import Dict, List, Optional, Type +from typing import Dict, List, Optional, Type, Union from google.protobuf.json_format import MessageToJson -from proto import Message +from google.protobuf.message import Message from feast.feature_view_projection import FeatureViewProjection from feast.field import Field +from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto +from feast.protos.feast.core.OnDemandFeatureView_pb2 import ( + OnDemandFeatureView as OnDemandFeatureViewProto, +) +from feast.protos.feast.core.StreamFeatureView_pb2 import ( + StreamFeatureView as StreamFeatureViewProto, +) class BaseFeatureView(ABC): @@ -89,7 +96,9 @@ def proto_class(self) -> Type[Message]: pass @abstractmethod - def to_proto(self) -> Message: + def to_proto( + self, + ) -> Union[FeatureViewProto, OnDemandFeatureViewProto, StreamFeatureViewProto]: pass @classmethod diff --git a/sdk/python/feast/infra/registry/base_registry.py b/sdk/python/feast/infra/registry/base_registry.py index c874001c89..c67164103e 100644 --- a/sdk/python/feast/infra/registry/base_registry.py +++ b/sdk/python/feast/infra/registry/base_registry.py @@ -19,7 +19,7 @@ from typing import Any, Dict, List, Optional from google.protobuf.json_format import MessageToJson -from proto import Message +from google.protobuf.message import Message from feast.base_feature_view import BaseFeatureView from feast.data_source import DataSource diff --git a/sdk/python/feast/infra/registry/registry.py b/sdk/python/feast/infra/registry/registry.py index 3c9843c904..d949b6079d 100644 --- a/sdk/python/feast/infra/registry/registry.py +++ b/sdk/python/feast/infra/registry/registry.py @@ -20,7 +20,7 @@ from urllib.parse import urlparse from google.protobuf.internal.containers import RepeatedCompositeFieldContainer -from proto import Message +from google.protobuf.message import Message from feast import usage from feast.base_feature_view import BaseFeatureView