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

27 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2023-02-04 02:28 +0100

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, f"_onnx_export_templates_{name}.tmpl") 

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

25 raise FileNotFoundError( # pragma: no cover 

26 f"Unable to find template {name!r} in folder {this!r}.") 

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

72 

73 

74def get_cpp_template(): 

75 """ 

76 Template to export :epkg:`ONNX` into a C++ code. 

77 """ 

78 return _get_file('cpp') 

79 

80 

81def get_python_template(): 

82 """ 

83 Template to export :epkg:`ONNX` into a python code. 

84 """ 

85 return _get_file('python')