def test_generate_entrypoint_python2(self):
    """ Test entrypoint generation for python2"""

    # prepare
    test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')

    # check
    generated_codes = _func_to_entrypoint(component_func=sample_component_func_two, python_version='python2')
    golden = '''\
def sample_component_func_two(a, b):
  result = 3.45
  if a == 'succ':
    result = float(b + 5)
  return result

def wrapper_sample_component_func_two(a,b,_output_file):
  output = sample_component_func_two(str(a),int(b))
  import os
  os.makedirs(os.path.dirname(_output_file))
  with open(_output_file, "w") as data:
    data.write(str(output))

import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())

if __name__ == "__main__":
  wrapper_sample_component_func_two(**args)
'''
    self.assertEqual(golden, generated_codes)
Example #2
0
    def test_func_to_entrypoint(self):
        """ Test entrypoint generation """

        # prepare
        test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')

        # check
        generated_codes = _func_to_entrypoint(
            component_func=sample_component_func)
        golden = '''\
def sample_component_func(a: str, b: int) -> float:
  result = 3.45
  if a == "succ":
    result = float(b + 5)
  return result

def wrapper_sample_component_func(a,b,_output_file):
  output = sample_component_func(str(a),int(b))
  import os
  os.makedirs(os.path.dirname(_output_file))
  with open(_output_file, "w") as data:
    data.write(str(output))

import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())

if __name__ == "__main__":
  wrapper_sample_component_func(**args)
'''
        self.assertEqual(golden, generated_codes)

        generated_codes = _func_to_entrypoint(
            component_func=sample_component_func_two)
        golden = '''\
def sample_component_func_two(a: str, b: int) -> float:
  result = 3.45
  if a == 'succ':
    result = float(b + 5)
  return result

def wrapper_sample_component_func_two(a,b,_output_file):
  output = sample_component_func_two(str(a),int(b))
  import os
  os.makedirs(os.path.dirname(_output_file))
  with open(_output_file, "w") as data:
    data.write(str(output))

import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())

if __name__ == "__main__":
  wrapper_sample_component_func_two(**args)
'''
        self.assertEqual(golden, generated_codes)

        generated_codes = _func_to_entrypoint(
            component_func=sample_component_func_three)
        golden = '''\
def sample_component_func_three() -> float:
  return 1.0

def wrapper_sample_component_func_three(_output_file):
  output = sample_component_func_three()
  import os
  os.makedirs(os.path.dirname(_output_file))
  with open(_output_file, "w") as data:
    data.write(str(output))

import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())

if __name__ == "__main__":
  wrapper_sample_component_func_three(**args)
'''
        self.assertEqual(golden, generated_codes)

        generated_codes = _func_to_entrypoint(
            component_func=sample_component_func_four)
        golden = '''\
from typing import NamedTuple
def sample_component_func_four() -> NamedTuple(
    'output', [('a', float), ('b', str)]):
  from collections import namedtuple
  output = namedtuple('output', ['a', 'b'])
  return output(1.0, 'test')

def wrapper_sample_component_func_four(_output_files):
  outputs = sample_component_func_four()
  import os
  for _output_file, output in zip(_output_files, outputs):
    os.makedirs(os.path.dirname(_output_file))
    with open(_output_file, "w") as data:
      data.write(str(output))

import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("_output_files", type=str, nargs=2)
args = vars(parser.parse_args())

if __name__ == "__main__":
  wrapper_sample_component_func_four(**args)
'''
        self.assertEqual(golden, generated_codes)