fix: Remove proto-plus imports (#4044)

tokoko · web-flow · commit ad8f5721af6d · 2024-03-27T12:53:55.000-05:00
remove proto-plus imports

Signed-off-by: tokoko &lt;togurg14@freeuni.edu.ge&gt;
diff --git 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
@@ -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
@@ -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
