more consistent value_type handling in Entity (feast-dev#1310)

oavdeev · web-flow · commit 2e5fb129a29c · 2021-02-02T21:40:41.000-08:00
Signed-off-by: Oleg Avdeev &lt;oleg.v.avdeev@gmail.com&gt;
diff --git a/infra/scripts/test-telemetry.sh b/infra/scripts/test-telemetry.sh
@@ -7,7 +7,6 @@ LOGS_ARTIFACT_PATH=/logs/artifacts
 
 pip install -r sdk/python/requirements-ci.txt
 make compile-protos-python
-make lint-python
 
 cd sdk/python/
 pip install -e .
diff --git a/sdk/python/feast/entity.py b/sdk/python/feast/entity.py
@@ -53,11 +53,6 @@ def __eq__(self, other):
         if not isinstance(other, Entity):
             raise TypeError("Comparisons should only involve Entity class objects.")
 
-        if isinstance(self.value_type, int):
-            self.value_type = ValueType(self.value_type).name
-        if isinstance(other.value_type, int):
-            other.value_type = ValueType(other.value_type).name
-
         if (
             self.labels != other.labels
             or self.name != other.name
@@ -100,7 +95,7 @@ def description(self, description):
         self._description = description
 
     @property
-    def value_type(self):
+    def value_type(self) -> ValueType:
         """
         Returns the type of this entity
         """
@@ -200,7 +195,7 @@ def from_proto(cls, entity_proto: EntityV2Proto):
         entity = cls(
             name=entity_proto.spec.name,
             description=entity_proto.spec.description,
-            value_type=ValueType(entity_proto.spec.value_type).name,  # type: ignore
+            value_type=ValueType(entity_proto.spec.value_type),
             labels=entity_proto.spec.labels,
         )
 
@@ -221,13 +216,11 @@ def to_proto(self) -> EntityV2Proto:
             created_timestamp=self.created_timestamp,
             last_updated_timestamp=self.last_updated_timestamp,
         )
-        if isinstance(self.value_type, ValueType):
-            self.value_type = self.value_type.value
 
         spec = EntitySpecProto(
             name=self.name,
             description=self.description,
-            value_type=self.value_type,
+            value_type=self.value_type.value,
             labels=self.labels,
         )
 
@@ -268,13 +261,10 @@ def to_spec_proto(self) -> EntitySpecProto:
             EntitySpecV2 protobuf
         """
 
-        if isinstance(self.value_type, ValueType):
-            self.value_type = self.value_type.value
-
         spec = EntitySpecProto(
             name=self.name,
             description=self.description,
-            value_type=self.value_type,
+            value_type=self.value_type.value,
             labels=self.labels,
         )
 
diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py
@@ -151,7 +151,7 @@ def _feature_table_to_argument(
         "project": project,
         "name": feature_table.name,
         "entities": [
-            {"name": n, "type": client.get_entity(n, project=project).value_type}
+            {"name": n, "type": client.get_entity(n, project=project).value_type.name}
             for n in feature_table.entities
         ],
         "max_age": feature_table.max_age.ToSeconds() if feature_table.max_age else None,
diff --git a/sdk/python/requirements-ci.txt b/sdk/python/requirements-ci.txt
@@ -10,7 +10,7 @@ mock==2.0.0
 pandavro==1.5.*
 moto
 mypy==0.790
-mypy-protobuf
+mypy-protobuf==1.24
 avro==1.10.0
 gcsfs
 urllib3>=1.25.4
diff --git a/sdk/python/requirements-dev.txt b/sdk/python/requirements-dev.txt
@@ -31,8 +31,8 @@ pandavro==1.5.*
 kafka-python==2.0.2
 tabulate==0.8.*
 isort>=5
-mypy
-mypy-protobuf
+mypy==0.790
+mypy-protobuf==1.24
 pre-commit
 flake8
 black==19.10b0
diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py
@@ -385,7 +385,7 @@ def test_apply_entity_success(self, test_client):
         assert (
             len(entities) == 1
             and entity.name == "driver_car_id"
-            and entity.value_type == ValueType(ValueProto.ValueType.STRING).name
+            and entity.value_type == ValueType(ValueProto.ValueType.STRING)
             and entity.description == "Car driver id"
             and "team" in entity.labels
             and entity.labels["team"] == "matchmaking"
