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)
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)