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
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-04 07:16 +0000
1from pathlib import Path
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
16class SourceFileGenerator:
17 """Context manager that controls the code generation process in generator scripts.
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.
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 """
34 def _scriptname(self) -> str:
35 import __main__
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 )
44 scriptpath = Path(__main__.__file__)
45 return scriptpath.name
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.")
55 scriptname = self._scriptname()
56 basename = scriptname.rsplit(".")[0]
58 from argparse import ArgumentParser
60 parser = ArgumentParser(
61 scriptname,
62 description="Generator script using pystencils-sfg",
63 allow_abbrev=False,
64 )
65 CommandLineParameters.add_args_to_parser(parser)
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 = []
73 cli_params = CommandLineParameters(sfg_args)
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)
80 self._header_only: bool = config.get_option("header_only")
81 self._output_dir: Path = config.get_option("output_directory")
83 output_files = config._get_output_files(basename)
85 from .ir import SfgSourceFile, SfgSourceFileType
87 self._header_file = SfgSourceFile(
88 output_files[0].name, SfgSourceFileType.HEADER
89 )
90 self._impl_file: SfgSourceFile | None
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))
100 # TODO: Find a way to not hard-code the restrict qualifier in pystencils
101 self._header_file.elements.append("#define RESTRICT __restrict__")
103 outer_namespace: str | _GlobalNamespace = config.get_option("outer_namespace")
105 namespace: str | None
106 if isinstance(outer_namespace, _GlobalNamespace):
107 namespace = None
108 else:
109 namespace = outer_namespace
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 )
121 sort_key = config.codestyle.get_option("includes_sorting_key")
122 if sort_key is None:
124 def default_key(h: HeaderFile):
125 return str(h)
127 sort_key = default_key
129 self._include_sort_key: Callable[[HeaderFile], Any] = sort_key
131 def clean_files(self):
132 header_path = self._output_dir / self._header_file.name
133 if header_path.exists():
134 header_path.unlink()
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()
141 def _finish_files(self) -> None:
142 from .ir import collect_includes
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)
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)
159 def _get_emitter(self):
160 return SfgCodeEmitter(
161 self._output_dir,
162 self._context.codestyle,
163 self._context.clang_format,
164 )
166 def __enter__(self) -> SfgComposer:
167 self.clean_files()
168 return SfgComposer(self._context)
170 def __exit__(self, exc_type, exc_value, traceback):
171 if exc_type is None:
172 self._finish_files()
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)