Coverage for src/pystencilssfg/generator.py: 95%

84 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-04 07:16 +0000

1from pathlib import Path 

2 

3from typing import Callable, Any 

4from .config import ( 

5 SfgConfig, 

6 CommandLineParameters, 

7 _GlobalNamespace, 

8) 

9from .context import SfgContext 

10from .composer import SfgComposer 

11from .emission import SfgCodeEmitter 

12from .exceptions import SfgException 

13from .lang import HeaderFile 

14 

15 

16class SourceFileGenerator: 

17 """Context manager that controls the code generation process in generator scripts. 

18 

19 The `SourceFileGenerator` must be used as a context manager by calling it within 

20 a ``with`` statement in the top-level code of a generator script. 

21 Upon entry to its context, it creates an `SfgComposer` which can be used to populate the generated files. 

22 When the managed region finishes, the code files are generated and written to disk at the locations 

23 defined by the configuration. 

24 Existing copies of the target files are deleted on entry to the managed region, 

25 and if an exception occurs within the managed region, no files are exported. 

26 

27 Args: 

28 sfg_config: Inline configuration for the code generator 

29 keep_unknown_argv: If `True`, any command line arguments given to the generator script 

30 that the `SourceFileGenerator` does not understand are stored in 

31 `sfg.context.argv <SfgContext.argv>`. 

32 """ 

33 

34 def _scriptname(self) -> str: 

35 import __main__ 

36 

37 if not hasattr(__main__, "__file__"): 

38 raise SfgException( 

39 "Invalid execution environment: " 

40 "It seems that you are trying to run the `SourceFileGenerator` in an environment " 

41 "without a valid entry point, such as a REPL or a multiprocessing fork." 

42 ) 

43 

44 scriptpath = Path(__main__.__file__) 

45 return scriptpath.name 

46 

47 def __init__( 

48 self, 

49 sfg_config: SfgConfig | None = None, 

50 keep_unknown_argv: bool = False, 

51 ): 

52 if sfg_config and not isinstance(sfg_config, SfgConfig): 

53 raise TypeError("sfg_config is not an SfgConfiguration.") 

54 

55 scriptname = self._scriptname() 

56 basename = scriptname.rsplit(".")[0] 

57 

58 from argparse import ArgumentParser 

59 

60 parser = ArgumentParser( 

61 scriptname, 

62 description="Generator script using pystencils-sfg", 

63 allow_abbrev=False, 

64 ) 

65 CommandLineParameters.add_args_to_parser(parser) 

66 

67 if keep_unknown_argv: 

68 sfg_args, script_args = parser.parse_known_args() 

69 else: 

70 sfg_args = parser.parse_args() 

71 script_args = [] 

72 

73 cli_params = CommandLineParameters(sfg_args) 

74 

75 config = cli_params.get_config() 

76 if sfg_config is not None: 

77 cli_params.find_conflicts(sfg_config) 

78 config.override(sfg_config) 

79 

80 self._header_only: bool = config.get_option("header_only") 

81 self._output_dir: Path = config.get_option("output_directory") 

82 

83 output_files = config._get_output_files(basename) 

84 

85 from .ir import SfgSourceFile, SfgSourceFileType 

86 

87 self._header_file = SfgSourceFile( 

88 output_files[0].name, SfgSourceFileType.HEADER 

89 ) 

90 self._impl_file: SfgSourceFile | None 

91 

92 if self._header_only: 

93 self._impl_file = None 

94 else: 

95 self._impl_file = SfgSourceFile( 

96 output_files[1].name, SfgSourceFileType.TRANSLATION_UNIT 

97 ) 

98 self._impl_file.includes.append(HeaderFile.parse(self._header_file.name)) 

99 

100 # TODO: Find a way to not hard-code the restrict qualifier in pystencils 

101 self._header_file.elements.append("#define RESTRICT __restrict__") 

102 

103 outer_namespace: str | _GlobalNamespace = config.get_option("outer_namespace") 

104 

105 namespace: str | None 

106 if isinstance(outer_namespace, _GlobalNamespace): 

107 namespace = None 

108 else: 

109 namespace = outer_namespace 

110 

111 self._context = SfgContext( 

112 self._header_file, 

113 self._impl_file, 

114 namespace, 

115 config.codestyle, 

116 config.clang_format, 

117 argv=script_args, 

118 project_info=cli_params.get_project_info(), 

119 ) 

120 

121 sort_key = config.codestyle.get_option("includes_sorting_key") 

122 if sort_key is None: 

123 

124 def default_key(h: HeaderFile): 

125 return str(h) 

126 

127 sort_key = default_key 

128 

129 self._include_sort_key: Callable[[HeaderFile], Any] = sort_key 

130 

131 def clean_files(self): 

132 header_path = self._output_dir / self._header_file.name 

133 if header_path.exists(): 

134 header_path.unlink() 

135 

136 if self._impl_file is not None: 

137 impl_path = self._output_dir / self._impl_file.name 

138 if impl_path.exists(): 

139 impl_path.unlink() 

140 

141 def _finish_files(self) -> None: 

142 from .ir import collect_includes 

143 

144 header_includes = collect_includes(self._header_file) 

145 self._header_file.includes = list( 

146 set(self._header_file.includes) | header_includes 

147 ) 

148 self._header_file.includes.sort(key=self._include_sort_key) 

149 

150 if self._impl_file is not None: 

151 impl_includes = collect_includes(self._impl_file) 

152 # If some header is already included by the generated header file, do not duplicate that inclusion 

153 impl_includes -= header_includes 

154 self._impl_file.includes = list( 

155 set(self._impl_file.includes) | impl_includes 

156 ) 

157 self._impl_file.includes.sort(key=self._include_sort_key) 

158 

159 def _get_emitter(self): 

160 return SfgCodeEmitter( 

161 self._output_dir, 

162 self._context.codestyle, 

163 self._context.clang_format, 

164 ) 

165 

166 def __enter__(self) -> SfgComposer: 

167 self.clean_files() 

168 return SfgComposer(self._context) 

169 

170 def __exit__(self, exc_type, exc_value, traceback): 

171 if exc_type is None: 

172 self._finish_files() 

173 

174 emitter = self._get_emitter() 

175 emitter.emit(self._header_file) 

176 if self._impl_file is not None: 

177 emitter.emit(self._impl_file)