Fixing DatasetKey.__new__ pytype, allowing object parameter type in case of a "flattened dataset key".

PiperOrigin-RevId: 506424291

Author: zoyahav

tensorflow_transform/beam/analyzer_cache.py

@@ -58,10 +58,12 @@ class DatasetKey(tfx_namedtuple.namedtuple('DatasetKey', ['key'])):
   __slots__ = ()
   _FLATTENED_DATASET_KEY = object()
 
-  def __new__(cls, dataset_key):
+  def __new__(cls, dataset_key: Union[str, object]) -> 'DatasetKey':
     if dataset_key is not DatasetKey._FLATTENED_DATASET_KEY:
-      # TODO(b/267425539): Remove pytype directive.
-      dataset_key = _make_valid_cache_component(dataset_key)  # pytype: disable=wrong-arg-types  # always-use-return-annotations
+      if not isinstance(dataset_key, str):
+        raise ValueError(
+            f'User provided dataset_key must be a str. Got: {dataset_key}')
+      dataset_key = _make_valid_cache_component(dataset_key)
     return super(DatasetKey, cls).__new__(cls, key=dataset_key)
 
   def __str__(self):

tensorflow_transform/beam/analyzer_cache_test.py

@@ -268,16 +268,17 @@ def write_to_file(value):
         analyzer_cache.DatasetKey('a'):
             analyzer_cache.DatasetCache({'b': [bytes([17, 19, 27, 31])]}, None)
     }
+    dataset_keys = list(test_cache_dict.keys())
 
     class LocalSource(beam.PTransform):
 
       def __init__(self, path):
         del path
 
       def expand(self, pbegin):
-        return pbegin | beam.Create([test_cache_dict['a'].cache_dict['b']])
+        return pbegin | beam.Create(
+            [test_cache_dict[k].cache_dict['b'] for k in dataset_keys])
 
-    dataset_keys = list(test_cache_dict.keys())
     cache_dir = self.get_temp_dir()
     with beam.Pipeline() as p:
       _ = test_cache_dict | analyzer_cache.WriteAnalysisCacheToFS(
@@ -289,9 +290,10 @@ def expand(self, pbegin):
       self.assertCountEqual(read_cache.keys(), ['a'])
       self.assertCountEqual(read_cache['a'].cache_dict.keys(), ['b'])
 
-      beam_test_util.assert_that(
-          read_cache['a'].cache_dict['b'],
-          beam_test_util.equal_to([test_cache_dict['a'].cache_dict['b']]))
+      for key in dataset_keys:
+        beam_test_util.assert_that(
+            read_cache[key].cache_dict['b'],
+            beam_test_util.equal_to([test_cache_dict[key].cache_dict['b']]))
 
 
 if __name__ == '__main__':