Coverage for mlprodict/onnx_tools/onnx_export_templates.py: 100%

1""" 

2@file 

3@brief Templates to export an ONNX graph in a way it can we created again 

4with a python script. 

5 

6.. versionadded:: 0.7 

7""" 

8import sys 

9import os 

10from textwrap import dedent 

11try: 

12 from functools import cache 

13except ImportError: # pragma: no cover 

14 # python 3.8 

15 from functools import lru_cache as cache 

16 

17 

18def _private_get_file(name): 

19 """ 

20 Retrieves one template. 

21 """ 

22 this = os.path.abspath(os.path.dirname(__file__)) 

23 filename = os.path.join(this, "_onnx_export_templates_%s.tmpl" % name) 

24 if not os.path.exists(filename): 

25 raise FileNotFoundError( # pragma: no cover 

26 "Unable to find template %r in folder %r." % (name, this)) 

27 with open(filename, "r", encoding="utf-8") as f: 

28 return dedent(f.read()) 

29 

30 

31if sys.version_info[:2] > (3, 6): 

32 @cache 

33 def _get_file(name): 

34 """ 

35 Retrieves one template. 

36 """ 

37 return _private_get_file(name) 

38else: # pragma: no cover 

39 def _get_file(name): 

40 """ 

41 Retrieves one template. 

42 """ 

43 return _private_get_file(name) 

44 

45 

46def get_onnx_template(): 

47 """ 

48 Template to export :epkg:`ONNX` into :epkg:`onnx` code. 

49 """ 

50 return _get_file('onnx') 

51 

52 

53def get_tf2onnx_template(): 

54 """ 

55 Template to export :epkg:`ONNX` into :epkg:`tensorflow-onnx` code. 

56 """ 

57 return _get_file('tf2onnx') 

58 

59 

60def get_numpy_template(): 

61 """ 

62 Template to export :epkg:`ONNX` into :epkg:`numpy` code. 

63 """ 

64 return _get_file('numpy') 

65 

66 

67def get_xop_template(): 

68 """ 

69 Template to export :epkg:`ONNX` into a code based on XOP API. 

70 """ 

71 return _get_file('xop')