gen_tasklet_reg.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. import argparse
  2. TEMPLATE = '''\
  3. {includes}\
  4. #include <tasklet/runtime/lib/{language}_wrapper.h>
  5. #include <tasklet/runtime/lib/registry.h>
  6. static const NTasklet::TRegHelper REG(
  7. "{name}",
  8. new NTasklet::{wrapper}
  9. );
  10. '''
  11. WRAPPER = {
  12. 'cpp': 'TCppWrapper<{impl}>()',
  13. 'js': 'TJsWrapper("{impl}")',
  14. 'go': 'TGoWrapper("{impl}")',
  15. 'py': 'TPythonWrapper("{impl}")',
  16. 'java': 'TJavaWrapper("{impl}", "{py_wrapper}")',
  17. }
  18. def parse_args():
  19. parser = argparse.ArgumentParser()
  20. parser.add_argument('name')
  21. parser.add_argument('output')
  22. parser.add_argument('-l', '--lang', choices=WRAPPER, required=True)
  23. parser.add_argument('-i', '--impl', required=True)
  24. parser.add_argument('-w', '--wrapper', required=False)
  25. parser.add_argument('includes', nargs='*')
  26. return parser.parse_args()
  27. if __name__ == '__main__':
  28. args = parse_args()
  29. includes = ''.join(
  30. '#include <{}>\n'.format(include)
  31. for include in args.includes
  32. )
  33. code = TEMPLATE.format(
  34. includes=includes,
  35. language=args.lang,
  36. name=args.name,
  37. wrapper=WRAPPER[args.lang].format(impl=args.impl, py_wrapper=args.wrapper),
  38. )
  39. with open(args.output, 'w') as f:
  40. f.write(code)