PY2 cleanup on the remaining tfx codebase.

PiperOrigin-RevId: 394170339

conftest.py

@@ -1,4 +1,3 @@
-# Lint as: python2, python3
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

package_build/ml-pipelines-sdk/package_config.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

package_build/tfx/package_config.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

package_config.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

setup.py

@@ -81,7 +81,7 @@ def finalize_options(self):
           'https://github.com/tensorflow/tfx/blob/master/package_build/'
           'README.md.\n\nEditable pip installation for development is still '
           'supported through `pip install -e`.')
-    super(_UnsupportedDevBuildWheelCommand, self).finalize_options()
+    super().finalize_options()
 
 
 class _BuildCommand(build.build):

tfx/__init__.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

tfx/benchmarks/__init__.py

@@ -1,4 +1,3 @@
-# Lint as: python2, python3
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

tfx/benchmarks/benchmark_base.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """Base class for benchmarks."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from absl import flags
 import apache_beam as beam
 from tensorflow.python.platform import test  # pylint: disable=g-direct-tensorflow-import

tfx/benchmarks/benchmark_dataset.py

@@ -13,15 +13,11 @@
 # limitations under the License.
 """Base class for classes representing a dataset for the benchmark."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import os
 import tensorflow as tf
 
 
-class BenchmarkDataset(object):
+class BenchmarkDataset:
   """Base class for classes representing a dataset for the benchmark."""
 
   def __init__(self, base_dir=None):

tfx/benchmarks/benchmark_utils.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """Utility functions shared across the different benchmarks."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import importlib
 from google.protobuf import text_format
 from tensorflow_metadata.proto.v0 import schema_pb2

tfx/benchmarks/datasets/__init__.py

@@ -1,4 +1,3 @@
-# Lint as: python2, python3
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

tfx/benchmarks/datasets/chicago_taxi/__init__.py

@@ -1,4 +1,3 @@
-# Lint as: python2, python3
 # Copyright 2019 Google LLC. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

tfx/benchmarks/datasets/chicago_taxi/dataset.py

@@ -13,16 +13,12 @@
 # limitations under the License.
 """Chicago taxi dataset."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import itertools
 import math
 import os
 import shutil
 import tempfile
-from typing import Optional, Text
+from typing import Optional
 
 from absl import logging
 import apache_beam as beam
@@ -216,8 +212,8 @@ class WideChicagoTaxiDataset(ChicagoTaxiDataset):
   _BUCKETIZE_KEYS = taxi_utils._BUCKET_FEATURE_KEYS
   _SCALE_KEYS = taxi_utils._DENSE_FLOAT_FEATURE_KEYS
 
-  def __init__(self, base_dir: Optional[Text] = None, num_analyzers: int = 10):
-    super(WideChicagoTaxiDataset, self).__init__(base_dir)
+  def __init__(self, base_dir: Optional[str] = None, num_analyzers: int = 10):
+    super().__init__(base_dir)
     self._num_vocabs = max(
         len(self._VOCABULARY_KEYS),
         math.ceil(num_analyzers * self._VOCABS_SHARE))

tfx/benchmarks/regenerate_datasets.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 """Tool to regenerate datasets used in benchmarks."""
 
-# Standard Imports
-
 from absl import app
 from absl import flags
 from absl import logging

tfx/benchmarks/tfma_benchmark_base.py

@@ -13,13 +13,8 @@
 # limitations under the License.
 """TFMA benchmark."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import time
 
-# Standard Imports
 
 import apache_beam as beam
 import tensorflow_model_analysis as tfma
@@ -36,14 +31,14 @@ class TFMABenchmarkBase(benchmark_base.BenchmarkBase):
   def __init__(self, dataset, **kwargs):
     # Benchmark runners may pass extraneous arguments we don't care about.
     del kwargs
-    super(TFMABenchmarkBase, self).__init__()
+    super().__init__()
     self._dataset = dataset
 
   def _max_num_examples(self):
     # TFMA is slower than TFT, so use a smaller number of examples from the
     # dataset.
     limit = 100000
-    parent_max = super(TFMABenchmarkBase, self)._max_num_examples()
+    parent_max = super()._max_num_examples()
     if parent_max is None:
       return limit
     return min(parent_max, limit)
@@ -57,7 +52,7 @@ def report_benchmark(self, **kwargs):
                                       getattr(tfx, "__version__", None))
     kwargs["extras"]["commit_tfma"] = (getattr(tfma, "GIT_COMMIT_ID", None) or
                                        getattr(tfma, "__version__", None))
-    super(TFMABenchmarkBase, self).report_benchmark(**kwargs)
+    super().report_benchmark(**kwargs)
 
   def benchmarkMiniPipeline(self):
     """Benchmark a "mini" version of TFMA - predict, slice and compute metrics.

tfx/benchmarks/tfma_benchmark_chicago_taxi.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """TFMA benchmark for Chicago Taxi dataset."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from tfx.benchmarks import tfma_benchmark_base
 from tfx.benchmarks.datasets.chicago_taxi import dataset
 
@@ -26,8 +22,7 @@
 class TFMABenchmarkChicagoTaxi(tfma_benchmark_base.TFMABenchmarkBase):
 
   def __init__(self, **kwargs):
-    super(TFMABenchmarkChicagoTaxi, self).__init__(
-        dataset=dataset.get_dataset(), **kwargs)
+    super().__init__(dataset=dataset.get_dataset(), **kwargs)
 
 
 if __name__ == "__main__":

tfx/benchmarks/tfma_v2_benchmark_base.py

@@ -13,14 +13,9 @@
 # limitations under the License.
 """TFMA v2 benchmark."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import copy
 import time
 
-# Standard Imports
 
 import apache_beam as beam
 import tensorflow as tf
@@ -60,7 +55,7 @@ class TFMAV2BenchmarkBase(benchmark_base.BenchmarkBase):
   def __init__(self, dataset, **kwargs):
     # Benchmark runners may pass extraneous arguments we don't care about.
     del kwargs
-    super(TFMAV2BenchmarkBase, self).__init__()
+    super().__init__()
     self._dataset = dataset
 
   def _init_model(self, multi_model, validation):
@@ -124,7 +119,7 @@ def _max_num_examples(self):
     # TFMA is slower than TFT, so use a smaller number of examples from the
     # dataset.
     limit = 100000
-    parent_max = super(TFMAV2BenchmarkBase, self)._max_num_examples()
+    parent_max = super()._max_num_examples()
     if parent_max is None:
       return limit
     return min(parent_max, limit)
@@ -143,7 +138,7 @@ def report_benchmark(self, **kwargs):
     # Stdout for use in tools which read the benchmark results from stdout.
     print(self._get_name(), kwargs["wall_time"],
           "({}x)".format(kwargs["iters"]))
-    super(TFMAV2BenchmarkBase, self).report_benchmark(**kwargs)
+    super().report_benchmark(**kwargs)
 
   def _runMiniPipeline(self, multi_model):
     """Benchmark a "mini" TFMA - predict, slice and compute metrics.

tfx/benchmarks/tfma_v2_benchmark_chicago_taxi.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """TFMA v2 benchmark for Chicago Taxi dataset."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from tfx.benchmarks import tfma_v2_benchmark_base
 from tfx.benchmarks.datasets.chicago_taxi import dataset
 
@@ -26,8 +22,7 @@
 class TFMAV2BenchmarkChicagoTaxi(tfma_v2_benchmark_base.TFMAV2BenchmarkBase):
 
   def __init__(self, **kwargs):
-    super(TFMAV2BenchmarkChicagoTaxi, self).__init__(
-        dataset=dataset.get_dataset(), **kwargs)
+    super().__init__(dataset=dataset.get_dataset(), **kwargs)
 
 
 if __name__ == "__main__":

tfx/benchmarks/tft_benchmark_base.py

@@ -13,16 +13,11 @@
 # limitations under the License.
 """TFT benchmark base."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import collections
 import shutil
 import tempfile
 import time
 
-# Standard Imports
 
 from absl import logging
 import apache_beam as beam
@@ -219,7 +214,7 @@ class TFTBenchmarkBase(benchmark_base.BenchmarkBase):
   def __init__(self, dataset, **kwargs):
     # Benchmark runners may pass extraneous arguments we don't care about.
     del kwargs
-    super(TFTBenchmarkBase, self).__init__()
+    super().__init__()
     self._dataset = dataset
 
   def report_benchmark(self, **kwargs):
@@ -231,7 +226,7 @@ def report_benchmark(self, **kwargs):
                                       getattr(tfx, "__version__", None))
     kwargs["extras"]["commit_tft"] = (getattr(tft, "GIT_COMMIT_ID", None) or
                                       getattr(tft, "__version__", None))
-    super(TFTBenchmarkBase, self).report_benchmark(**kwargs)
+    super().report_benchmark(**kwargs)
 
   def _benchmarkAnalyzeAndTransformDatasetCommon(self, force_tf_compat_v1):
     """Common implementation to benchmark AnalyzeAndTransformDataset."""

tfx/benchmarks/tft_benchmark_chicago_taxi.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """TFT benchmark for Chicago Taxi dataset."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from absl import flags
 from tfx.benchmarks import tft_benchmark_base
 from tfx.benchmarks.datasets.chicago_taxi import dataset
@@ -32,14 +28,13 @@
 class TFTBenchmarkChicagoTaxi(tft_benchmark_base.TFTBenchmarkBase):
 
   def __init__(self, **kwargs):
-    super(TFTBenchmarkChicagoTaxi, self).__init__(
-        dataset=dataset.get_dataset(), **kwargs)
+    super().__init__(dataset=dataset.get_dataset(), **kwargs)
 
 
 class TFTBenchmarkChicagoTaxiWide(tft_benchmark_base.TFTBenchmarkBase):
 
   def __init__(self, **kwargs):
-    super(TFTBenchmarkChicagoTaxiWide, self).__init__(
+    super().__init__(
         dataset=dataset.get_wide_dataset(num_analyzers=self._num_analyzers()),
         **kwargs)
 

tfx/dsl/components/base/base_component.py

@@ -15,7 +15,7 @@
 
 import abc
 import inspect
-from typing import Any, Dict, Optional, Text, Union
+from typing import Any, Dict, Optional, Union
 
 from tfx import types
 from tfx.dsl.components.base import base_driver
@@ -161,11 +161,11 @@ def __repr__(self):
 
   @property
   @doc_controls.do_not_doc_in_subclasses
-  def inputs(self) -> Dict[Text, Any]:
+  def inputs(self) -> Dict[str, Any]:
     return self.spec.inputs
 
   @property
-  def outputs(self) -> Dict[Text, Any]:
+  def outputs(self) -> Dict[str, Any]:
     """Component's output channel dict."""
     return self.spec.outputs
 

tfx/dsl/components/base/base_node.py

@@ -14,7 +14,7 @@
 """Base class for TFX nodes."""
 
 import abc
-from typing import Any, Dict, Optional, Text, Type
+from typing import Any, Dict, Optional, Type
 
 from tfx.dsl.components.base import base_driver
 from tfx.dsl.components.base import base_executor
@@ -122,12 +122,12 @@ def with_id(self, id: str) -> 'BaseNode':  # pylint: disable=redefined-builtin
 
   @property
   @abc.abstractmethod
-  def inputs(self) -> Dict[Text, Any]:
+  def inputs(self) -> Dict[str, Any]:
     pass
 
   @property
   @abc.abstractmethod
-  def outputs(self) -> Dict[Text, Any]:
+  def outputs(self) -> Dict[str, Any]:
     pass
 
   @property

tfx/dsl/components/base/node_registry_test.py

@@ -14,7 +14,7 @@
 """Tests for tfx.dsl.components.base.node_registry."""
 
 import threading
-from typing import Any, Dict, Text
+from typing import Any, Dict
 
 import tensorflow as tf
 from tfx.dsl.components.base import base_node
@@ -24,11 +24,11 @@
 class _FakeNode(base_node.BaseNode):
 
   @property
-  def inputs(self) -> Dict[Text, Any]:
+  def inputs(self) -> Dict[str, Any]:
     return {}
 
   @property
-  def outputs(self) -> Dict[Text, Any]:
+  def outputs(self) -> Dict[str, Any]:
     return {}
 
   @property