Use to_flyte_literal_type() when comparing schema columns (#128)

flytekit/__init__.py

@@ -2,4 +2,4 @@
 
 import flytekit.plugins
 
-__version__ = '0.9.3'
+__version__ = '0.9.4'

flytekit/common/types/impl/schema.py

@@ -947,7 +947,8 @@ def cast_to(self, other_type):
  additional_msg="Cannot cast because a required column '{}' was not found.".format(k),
  received_value=self
  )
- if v != self.type.sdk_columns[k]:
+ if not isinstance(v, _base_sdk_types.FlyteSdkType) or \
+ v.to_flyte_literal_type() != self.type.sdk_columns[k].to_flyte_literal_type():
  raise _user_exceptions.FlyteTypeException(
  self.type.sdk_columns[k],
  v,

tests/flytekit/unit/common_tests/types/test_schema.py

@@ -52,3 +52,20 @@ def test_typed_schema():
  assert len(b.type.columns) == len(_ALL_COLUMN_TYPES)
  assert list(b.type.sdk_columns.items()) == _ALL_COLUMN_TYPES
  assert b.remote_location.startswith(t.name)
+
+
+# Ensures that subclassing types works inside a schema.
+def test_casting():
+ class MyDateTime(primitives.Datetime):
+ ...
+
+ with test_utils.LocalTestFileSystem() as t:
+ test_columns_1 = [('altered', MyDateTime)]
+ test_columns_2 = [('altered', primitives.Datetime)]
+
+ instantiator_1 = schema.schema_instantiator(test_columns_1)
+ a = instantiator_1()
+
+ instantiator_2 = schema.schema_instantiator(test_columns_2)
+
+ a.cast_to(instantiator_2._schema_type)