Fix lint errors

tfx/components/example_diff/executor.py

@@ -175,7 +175,11 @@ def Do(self, input_dict: Dict[str, List[types.Artifact]],
         logging.info('Processing split pair %s', split_pair)
         # pylint: disable=cell-var-from-loop
         @beam.ptransform_fn
-        def _iteration(p):
+        def _iteration(p,
+                       test_tfxio=test_tfxio,
+                       base_tfxio=base_tfxio,
+                       split_pair=split_pair
+        ):
           base_examples = (
               p | 'TFXIORead[base]' >> test_tfxio.RawRecordBeamSource()
               | 'Parse[base]' >> beam.Map(_parse_example))

tfx/components/example_gen/utils.py

@@ -271,8 +271,8 @@ def _make_zero_padding_spec_value(spec_full_regex: str, pattern: str,
     width_int = 0
     try:
       width_int = int(width_str)
-    except ValueError:
-      raise ValueError('Width modifier is not a integer: %s' % pattern)
+    except ValueError as e:
+      raise ValueError('Width modifier is not a integer: %s' % pattern) from e
     if width_int <= 0:
       raise ValueError('Width modifier is not positive: %s' % pattern)
     if width_int < len(str(spec_value)):
@@ -377,30 +377,34 @@ def _find_matched_span_version_from_path(
     matched_span_tokens = [result.group(SPAN_PROPERTY_NAME)]
     try:
       matched_span_int = int(matched_span_tokens[0])
-    except ValueError:
+    except ValueError as e:
       raise ValueError('Cannot find %s number from %s based on %s' %
-                       (SPAN_PROPERTY_NAME, file_path, split_regex_pattern))
+                       (SPAN_PROPERTY_NAME, file_path, split_regex_pattern)
+                       ) from e
   elif is_match_date:
     matched_span_tokens = [
         result.group(name) for name in ['year', 'month', 'day']
     ]
     try:
       matched_span_ints = [int(elem) for elem in matched_span_tokens]
-    except ValueError:
+    except ValueError as e:
       raise ValueError('Cannot find %s number using date from %s based on %s' %
-                       (SPAN_PROPERTY_NAME, file_path, split_regex_pattern))
+                       (SPAN_PROPERTY_NAME, file_path, split_regex_pattern)
+                       ) from e
     try:
       matched_span_int = date_to_span_number(*matched_span_ints)
-    except ValueError:
-      raise ValueError('Retrieved date is invalid for file: %s' % file_path)
+    except ValueError as e:
+      raise ValueError('Retrieved date is invalid for file: %s' % file_path
+                       ) from e
 
   if is_match_version:
     matched_version = result.group(VERSION_PROPERTY_NAME)
     try:
       matched_version_int = int(matched_version)
-    except ValueError:
+    except ValueError as e:
       raise ValueError('Cannot find %s number from %s based on %s' %
-                       (VERSION_PROPERTY_NAME, file_path, split_regex_pattern))
+                       (VERSION_PROPERTY_NAME, file_path, split_regex_pattern)
+                       ) from e
 
   return (matched_span_tokens, matched_span_int, matched_version,
           matched_version_int)
@@ -417,10 +421,10 @@ def _get_spec_width(spec_full_regex: str, spec_name: str,
       width_int = int(width_str)
       if width_int <= 0:
         raise ValueError('Not a positive integer.')
-    except ValueError:
+    except ValueError as e:
       raise ValueError(
           'Width modifier in %s spec is not a positive integer: %s' %
-          (spec_name, split.pattern))
+          (spec_name, split.pattern)) from e
   return width_str
 
 

tfx/components/trainer/rewriting/rewriter.py

@@ -100,18 +100,18 @@ def perform_rewrite(self, original_model: ModelDescription,
       raise ValueError('{} failed to perform pre-rewrite validation. Original '
                        'model: {}. Error: {}'.format(self.name,
                                                      str(original_model),
-                                                     str(v)))
+                                                     str(v))) from v
 
     try:
       self._rewrite(original_model, rewritten_model)
     except ValueError as v:
       raise ValueError(
           '{} failed to rewrite model. Original model: {}. Error {}'.format(
-              self.name, str(original_model), str(v)))
+              self.name, str(original_model), str(v))) from v
 
     try:
       self._post_rewrite_validate(rewritten_model)
     except ValueError as v:
       raise ValueError(
           '{} failed to validate rewritten model. Rewritten model: {}. Error {}'
-          .format(self.name, str(rewritten_model), str(v)))
+          .format(self.name, str(rewritten_model), str(v))) from v

tfx/dsl/compiler/node_inputs_compiler_test.py

@@ -47,7 +47,7 @@ class DummyArtifact(types.Artifact):
 
 class DummyNode(base_node.BaseNode):
 
-  def __init__(self, id: str, inputs=None, exec_properties=None):  # pylint: disable=redefined-builtin
+  def __init__(self, id: str, inputs=None, exec_properties=None):  # noqa: A002
     super().__init__()
     self.with_id(id)
     self._inputs = inputs or {}

tfx/dsl/compiler/placeholder_utils.py

@@ -391,9 +391,9 @@ def _resolve_property_operator(
       if op.is_custom_property:
         return value.get_custom_property(op.key)
       return value.__getattr__(op.key)
-    except:
+    except Exception as e:
       raise ValueError("ArtifactPropertyOperator failed to find property with "
-                       f"key {op.key}.")
+                       f"key {op.key}.") from e
 
   @_register(placeholder_pb2.Base64EncodeOperator)
   def _resolve_base64_encode_operator(
@@ -514,18 +514,18 @@ def _resolve_proto_operator(
         if field.startswith("."):
           try:
             value = getattr(value, field[1:])
-          except AttributeError:
+          except AttributeError as e:
             raise ValueError("While evaluting placeholder proto operator, "
                              f"got unknown proto field {field} on proto of "
-                             f"type {type(value)}.")
+                             f"type {type(value)}.") from e
           continue
         map_key = re.findall(r"\[['\"](.+)['\"]\]", field)
         if len(map_key) == 1:
           try:
             value = value[map_key[0]]
-          except KeyError:
+          except KeyError as e:
             raise ValueError("While evaluting placeholder proto operator, "
-                             f"got unknown map field {field}.")
+                             f"got unknown map field {field}.") from e
           continue
         # Going forward, index access for proto fields should be handled by
         # index op. This code here is kept to avoid breaking existing executor
@@ -534,9 +534,9 @@ def _resolve_proto_operator(
         if index and str.isdecimal(index[0]):
           try:
             value = value[int(index[0])]
-          except IndexError:
+          except IndexError as e:
             raise ValueError("While evaluting placeholder proto operator, "
-                             f"got unknown index field {field}.")
+                             f"got unknown index field {field}.") from e
           continue
         raise ValueError(f"Got unsupported proto field path: {field}")
 

tfx/dsl/component/experimental/decorators.py

@@ -202,7 +202,7 @@ def __init__(self, *unused_args, **kwargs):
       json_compat_typehint = getattr(channel_parameter, '_JSON_COMPAT_TYPEHINT',
                                      None)
       if json_compat_typehint:
-        setattr(spec_kwargs[key], '_JSON_COMPAT_TYPEHINT', json_compat_typehint)
+        spec_kwargs[key]._JSON_COMPAT_TYPEHINT = json_compat_typehint
     spec = self.SPEC_CLASS(**spec_kwargs)
     super().__init__(spec)
     # Set class name, which is the decorated function name, as the default id.

tfx/dsl/component/experimental/decorators_test.py

@@ -418,7 +418,7 @@ def testDefinitionInClosureFails(self):
         'the module level',
     ):
       @component
-      def my_component():  # pylint: disable=unused-variable
+      def my_component():  # noqa: F841
         return None
 
   def testNonKwargFails(self):

tfx/dsl/component/experimental/decorators_typeddict_test.py

@@ -432,7 +432,7 @@ def testDefinitionInClosureFails(self):
     ):
 
       @component
-      def my_component():  # pylint: disable=unused-variable
+      def my_component():  # noqa: F841
         return None
 
   def testNonKwargFails(self):

tfx/dsl/component/experimental/utils.py

@@ -209,21 +209,13 @@ def _create_component_spec_class(
           type=artifact_type, optional=(key in arg_defaults)
       )
       if json_compatible_inputs and key in json_compatible_inputs:
-        setattr(
-            spec_inputs[key],
-            '_JSON_COMPAT_TYPEHINT',
-            json_compatible_inputs[key],
-        )
+        spec_inputs[key]._JSON_COMPAT_TYPEHINT = json_compatible_inputs[key]
   if outputs:
     for key, artifact_type in outputs.items():
       assert key not in arg_defaults, 'Optional outputs are not supported.'
       spec_outputs[key] = component_spec.ChannelParameter(type=artifact_type)
       if json_compatible_outputs and key in json_compatible_outputs:
-        setattr(
-            spec_outputs[key],
-            '_JSON_COMPAT_TYPEHINT',
-            json_compatible_outputs[key],
-        )
+        spec_outputs[key]._JSON_COMPAT_TYPEHINT = json_compatible_outputs[key]
   if parameters:
     for key, param_type in parameters.items():
       if inspect.isclass(param_type) and issubclass(

tfx/dsl/components/base/base_component_test.py

@@ -61,7 +61,7 @@ class _BasicComponent(base_component.BaseComponent):
   def __init__(self,
                spec: types.ComponentSpec = None,
                folds: int = None,
-               input: types.Channel = None):  # pylint: disable=redefined-builtin
+               input: types.Channel = None):  # noqa: A002
     if not spec:
       output = types.Channel(type=_OutputArtifact)
       spec = _BasicComponentSpec(folds=folds, input=input, output=output)

tfx/dsl/components/base/base_node.py

@@ -126,12 +126,12 @@ def component_id(self) -> str:
 
   @id.setter
   @doc_controls.do_not_doc_in_subclasses
-  def id(self, id: str) -> None:  # pylint: disable=redefined-builtin
+  def id(self, id: str) -> None:  # noqa: A002
     self._id = id
 
   # TODO(kmonte): Update this to Self once we're on 3.11 everywhere
   @doc_controls.do_not_doc_in_subclasses
-  def with_id(self, id: str) -> typing_extensions.Self:  # pylint: disable=redefined-builtin
+  def with_id(self, id: str) -> typing_extensions.Self:  # noqa: A002
     self._id = id
     return self
 

tfx/dsl/control_flow/for_each_test.py

@@ -88,8 +88,8 @@ def testForEach_LoopVariableNotUsed_Disallowed(self):
     with self.subTest('Not using at all'):
       with self.assertRaises(ValueError):
         a = A()
-        with for_each.ForEach(a.outputs['aa']) as aa:  # pylint: disable=unused-variable
-          b = B()  # aa is not used. # pylint: disable=unused-variable
+        with for_each.ForEach(a.outputs['aa']) as aa:  # noqa: F841
+          b = B()  # aa is not used. # noqa: F841
 
     with self.subTest('Source channel is not a loop variable.'):
       with self.assertRaises(ValueError):
@@ -116,9 +116,9 @@ def testForEach_DifferentLoop_HasDifferentContext(self):
     a = A()
     b = B()
     with for_each.ForEach(a.outputs['aa']) as aa1:
-      c1 = C(aa=aa1, bb=b.outputs['bb'])  # pylint: disable=unused-variable
+      c1 = C(aa=aa1, bb=b.outputs['bb'])  # noqa: F841
     with for_each.ForEach(a.outputs['aa']) as aa2:
-      c2 = C(aa=aa2, bb=b.outputs['bb'])  # pylint: disable=unused-variable
+      c2 = C(aa=aa2, bb=b.outputs['bb'])  # noqa: F841
 
     context1 = dsl_context_registry.get().get_contexts(c1)[-1]
     context2 = dsl_context_registry.get().get_contexts(c2)[-1]

tfx/dsl/input_resolution/canned_resolver_functions.py

@@ -744,7 +744,7 @@ def pick(channel: channel_types.BaseChannel, i: int, /):
   return _slice(channel, start=i, stop=(i + 1) or None, min_count=1)
 
 
-def slice(  # pylint: disable=redefined-builtin
+def slice(  # noqa: A002, A001
     channel: channel_types.BaseChannel,
     /,
     start: Optional[int] = None,

tfx/dsl/input_resolution/ops/test_utils.py

@@ -51,10 +51,10 @@ class DummyArtifact(types.Artifact):
       'version': tfx_artifact.Property(type=tfx_artifact.PropertyType.INT),
   }
 
-  # pylint: disable=redefined-builtin
+  # noqa: A002
   def __init__(
       self,
-      id: Optional[str] = None,
+      id: Optional[str] = None, # noqa: A001, A002
       uri: Optional[str] = None,
       create_time_since_epoch: Optional[int] = None,
   ):
@@ -149,7 +149,7 @@ class FakeComponent(base_component.BaseComponent):
 
   DRIVER_CLASS = base_driver.BaseDriver
 
-  def __init__(self, id: str, inputs=None, exec_properties=None):  # pylint: disable=redefined-builtin
+  def __init__(self, id: str, inputs=None, exec_properties=None):  # noqa: A002
     super().__init__(spec=FakeSpec())
     self.with_id(id)
 

tfx/dsl/input_resolution/resolver_op.py

@@ -228,7 +228,7 @@ def my_resolver_stem(input_dict):
 
   def __init__(
       self, *,
-      type: Type[_T],  # pylint: disable=redefined-builtin
+      type: Type[_T],  # noqa: A002
       default: Union[_T, _Empty] = _EMPTY):
     self._type = type
     self._required = default is _EMPTY

tfx/dsl/input_resolution/resolver_op_test.py

@@ -104,7 +104,7 @@ class ResolverOpTest(tf.test.TestCase):
 
   def testDefineOp_PropertyDefaultViolatesType(self):
     with self.assertRaises(TypeError):
-      class BadProperty(resolver_op.ResolverOp):  # pylint: disable=unused-variable
+      class BadProperty(resolver_op.ResolverOp):  # noqa: F841
         str_prop = resolver_op.Property(type=str, default=42)
 
   def testOpCall_ReturnsOpNode(self):

tfx/dsl/io/fileio.py

@@ -20,8 +20,8 @@
 from tfx.dsl.io.filesystem import PathType
 
 # Import modules that may provide filesystem plugins.
-import tfx.dsl.io.plugins.tensorflow_gfile  # pylint: disable=unused-import, g-import-not-at-top
-import tfx.dsl.io.plugins.local  # pylint: disable=unused-import, g-import-not-at-top
+import tfx.dsl.io.plugins.tensorflow_gfile  # noqa: F401, E402
+import tfx.dsl.io.plugins.local  # noqa: F401, E402
 
 
 # Expose `NotFoundError` as `fileio.NotFoundError`.
@@ -33,7 +33,7 @@ def _get_filesystem(path) -> Type[filesystem.Filesystem]:
           .get_filesystem_for_path(path))
 
 
-def open(path: PathType, mode: str = 'r'):  # pylint: disable=redefined-builtin
+def open(path: PathType, mode: str = 'r'):  # noqa: A002, A001
   """Open a file at the given path."""
   return _get_filesystem(path).open(path, mode=mode)
 

tfx/dsl/placeholder/artifact_placeholder.py

@@ -23,7 +23,7 @@
 _types = placeholder_base.types
 
 
-def input(key: str) -> ArtifactPlaceholder:  # pylint: disable=redefined-builtin
+def input(key: str) -> ArtifactPlaceholder:  # noqa: A002, A001
   """Returns a Placeholder that represents an input artifact.
 
   Args:

tfx/dsl/placeholder/placeholder_base.py

@@ -45,7 +45,8 @@
 # where ComponentSpecIntf is an interface that gives access to the few things
 # that Placeholder::encode() needs from the spec. TODO(pke) Implement.
 # TODO(b/191610358): Reduce the number of circular type-dependencies.
-types = Any  # To resolve circular dependency caused by type annotations.
+# To resolve circular dependency caused by type annotations.
+types = Any  # noqa: F811
 
 # TODO(b/190409099): Support RuntimeParameter.
 ValueType = Union[int, float, str, bool]

tfx/dsl/placeholder/proto_placeholder.py

@@ -509,11 +509,11 @@ def _purge_types(
         _purge_types(name_prefix, message_descriptor)  # Step 2
       _remove_unless(
           file_descriptor.message_type,
-          lambda m: f'{name_prefix}.{m.name}' in self._keep_types,  # pylint: disable=cell-var-from-loop
+          lambda m, name_prefix=name_prefix: f'{name_prefix}.{m.name}' in self._keep_types,  # pylint: disable=cell-var-from-loop
       )
       _remove_unless(
           file_descriptor.enum_type,
-          lambda e: f'{name_prefix}.{e.name}' in self._keep_types,  # pylint: disable=cell-var-from-loop
+          lambda e, name_prefix=name_prefix: f'{name_prefix}.{e.name}' in self._keep_types,  # pylint: disable=cell-var-from-loop
       )
 
     # Step 4: Remove file descriptors that became empty. Remove declared

tfx/examples/bert/utils/bert_models.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Configurable fine-tuning BERT models for various tasks."""
 
+from __future__ import annotations
 from typing import Optional, List, Union
 
 import tensorflow as tf
@@ -59,8 +60,7 @@ def build_bert_classifier(bert_layer: tf.keras.layers.Layer,
 
 def compile_bert_classifier(
     model: tf.keras.Model,
-    loss: tf.keras.losses.Loss = tf.keras.losses.SparseCategoricalCrossentropy(
-        from_logits=True),
+    loss: tf.keras.losses.Loss | None = None,
     learning_rate: float = 2e-5,
     metrics: Optional[List[Union[str, tf.keras.metrics.Metric]]] = None):
   """Compile the BERT classifier using suggested parameters.
@@ -79,6 +79,9 @@ def compile_bert_classifier(
   Returns:
     None.
   """
+  if loss is None:
+    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+
   if metrics is None:
     metrics = ["sparse_categorical_accuracy"]