# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
# Lint as: python3
"""Executor for Transform."""

import inspect
from typing import Any, Mapping, Text

from tfx import components
from tfx.components.transform import executor
from tfx.components.transform import labels


class Executor(executor.Executor):
  """Transform executor."""

  def _GetPreprocessingFn(self, inputs: Mapping[Text, Any],
                          unused_outputs: Mapping[Text, Any]) -> Any:
    """See base class."""
    fn = super(Executor, self)._GetPreprocessingFn(inputs, unused_outputs)
    # The internal function's inputs argument shadows the external inputs
    # argument. We can't rename either in case inputs are specified as kwargs,
    # so we keep a reference to the external inputs here instead.
    inputs_dict = inputs

    def preprocessing_fn(inputs):
      """The preprocessing_fn to return."""
      args = {}
      argspec = inspect.getfullargspec(fn).args
      if 'schema' in argspec:
        schema_path = components.util.value_utils.GetSoleValue(
            inputs_dict, labels.SCHEMA_PATH_LABEL)
        schema = self._GetSchema(schema_path)
        args['schema'] = schema
      return fn(inputs, **args)

    return preprocessing_fn

@@ -1021,23 +1021,16 @@
         
-unused_
 outputs:
@@ -1147,15 +1147,8 @@
 ts, 
-unused_
 outp
@@ -1477,24 +1477,26 @@
 n.%22%22%22%0A      
+kw
 args = %7B%7D%0A  
@@ -1541,16 +1541,17 @@
 n).args%0A
+%0A
       if
@@ -1737,24 +1737,26 @@
 th)%0A        
+kw
 args%5B'schema
@@ -1766,16 +1766,225 @@
 = schema
+%0A%0A      if 'transform_graph_dir' in argspec:%0A        kwargs%5B%0A            'transform_graph_dir'%5D = components.util.value_utils.GetSoleValue(%0A                outputs, labels.TRANSFORM_METADATA_OUTPUT_PATH_LABEL)
 %0A      r
@@ -2002,16 +2002,18 @@
 puts, **
+kw
 args)%0A%0A 

Update NitroML's custom Transform component to pass the TransforOuputDir to the `preprocessing_fn`.