diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1d0cdda1b64ec9eff93367589fe3a08d975deb7
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,11 @@
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def prepare_composer(doctest_namespace):
+ from pystencilssfg import SfgContext, SfgComposer
+
+ # Place a composer object in the environment for doctests
+
+ sfg = SfgComposer(SfgContext())
+ doctest_namespace["sfg"] = sfg
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 84c2b779b553fe22fa8ef23a2f0a4872c9edb57c..11d64fd5a142eab3f1558c52944b772f94a5f266 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -27,6 +27,7 @@ extensions = [
"myst_parser",
"sphinx.ext.autodoc",
"sphinx.ext.napoleon",
+ "sphinx.ext.doctest",
"sphinx.ext.intersphinx",
"sphinx_autodoc_typehints",
"sphinx_design",
@@ -64,6 +65,13 @@ intersphinx_mapping = {
autodoc_member_order = "bysource"
autodoc_typehints = "description"
+# Doctest Setup
+
+doctest_global_setup = '''
+from pystencilssfg import SfgContext, SfgComposer
+sfg = SfgComposer(SfgContext())
+'''
+
# Prepare code generation examples
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000000000000000000000000000000000000..94a3a6c5cd76967b094d0e26bab983291057461d
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,8 @@
+[pytest]
+testpaths = src/pystencilssfg tests/
+python_files = "test_*.py"
+# Need to ignore the generator scripts, otherwise they would be executed
+# during test collection
+addopts = --doctest-modules --ignore=tests/generator_scripts/scripts
+
+doctest_optionflags = NORMALIZE_WHITESPACE IGNORE_EXCEPTION_DETAIL
diff --git a/src/pystencilssfg/composer/basic_composer.py b/src/pystencilssfg/composer/basic_composer.py
index 7fda6c9ae2b115bbafdc40ee670625ceec61a578..6a2472067ddb746a4d0700599c1541491d86e4a9 100644
--- a/src/pystencilssfg/composer/basic_composer.py
+++ b/src/pystencilssfg/composer/basic_composer.py
@@ -280,6 +280,18 @@ class SfgBasicComposer(SfgIComposer):
"""Create a variable with given name and data type."""
return AugExpr(create_type(dtype)).var(name)
+ def vars(self, names: str, dtype: UserTypeSpec) -> tuple[AugExpr, ...]:
+ """Create multiple variables with given names and the same data type.
+
+ Example:
+
+ >>> sfg.vars("x, y, z", "float32")
+ (x, y, z)
+
+ """
+ varnames = names.split(",")
+ return tuple(self.var(n.strip(), dtype) for n in varnames)
+
def init(self, lhs: VarLike) -> SfgInplaceInitBuilder:
"""Create a C++ in-place initialization.
@@ -298,10 +310,37 @@ class SfgBasicComposer(SfgIComposer):
"""
return SfgInplaceInitBuilder(_asvar(lhs))
- def expr(self, fmt: str, *deps, **kwdeps):
+ def expr(self, fmt: str, *deps, **kwdeps) -> AugExpr:
"""Create an expression while keeping track of variables it depends on.
-
-
+
+ This method is meant to be used similarly to `str.format`; in fact,
+ it calls `str.format` internally and therefore supports all of its
+ formatting features.
+ In addition, however, the format arguments are scanned for *variables*
+ (e.g. created using `var`), which are attached to the expression.
+ This way, *pystencils-sfg* keeps track of any variables an expression depends on.
+
+ :Example:
+
+ >>> x, y, z, w = sfg.vars("x, y, z, w", "float32")
+ >>> expr = sfg.expr("{} + {} * {}", x, y, z)
+ >>> expr
+ x + y * z
+
+ You can look at the expression's dependencies:
+
+ >>> sorted(expr.depends, key=lambda v: v.name)
+ [x: float, y: float, z: float]
+
+ If you use an existing expression to create a larger one, the new expression
+ inherits all variables from its parts:
+
+ >>> expr2 = sfg.expr("{} + {}", expr, w)
+ >>> expr2
+ x + y * z + w
+ >>> sorted(expr2.depends, key=lambda v: v.name)
+ [w: float, x: float, y: float, z: float]
+
"""
return AugExpr.format(fmt, *deps, **kwdeps)
@@ -338,7 +377,9 @@ class SfgBasicComposer(SfgIComposer):
def set_param(self, param: VarLike | sp.Symbol, expr: ExprLike):
depends = _depends(expr)
- var = _asvar(param) if isinstance(param, _VarLike) else param
+ var: SfgVar | sp.Symbol = (
+ _asvar(param) if isinstance(param, _VarLike) else param
+ )
return SfgDeferredParamSetter(var, depends, str(expr))
def map_param(
@@ -351,7 +392,9 @@ class SfgBasicComposer(SfgIComposer):
side object from one or multiple right-hand side dependencies."""
if isinstance(depends, _VarLike):
depends = [depends]
- lhs_var = _asvar(param) if isinstance(param, _VarLike) else param
+ lhs_var: SfgVar | sp.Symbol = (
+ _asvar(param) if isinstance(param, _VarLike) else param
+ )
return SfgDeferredParamMapping(
lhs_var, set(_asvar(v) for v in depends), mapping
)
@@ -363,7 +406,7 @@ class SfgBasicComposer(SfgIComposer):
lhs_components: Vector components as a list of symbols.
rhs: A `SrcVector` object representing a vector data structure.
"""
- components = [
+ components: list[SfgVar | sp.Symbol] = [
(_asvar(c) if isinstance(c, _VarLike) else c) for c in lhs_components
]
return SfgDeferredVectorMapping(components, rhs)
diff --git a/src/pystencilssfg/ir/source_components.py b/src/pystencilssfg/ir/source_components.py
index 2eab9935ae7cc5be7fa2d6a766ed22bd2ee121dd..859f9266de18fb952405721892136bde3fda3fd9 100644
--- a/src/pystencilssfg/ir/source_components.py
+++ b/src/pystencilssfg/ir/source_components.py
@@ -250,7 +250,7 @@ class SfgVar:
return self._name
def __repr__(self) -> str:
- return f"SfgVar( {self._name}, {repr(self._dtype)} )"
+ return f"{self._name}: {self._dtype}"
SymbolLike_T = TypeVar("SymbolLike_T", bound=KernelParameter)
diff --git a/src/pystencilssfg/lang/expressions.py b/src/pystencilssfg/lang/expressions.py
index ca5a78c4d0794d33a2b22827b7a5e9d5a09e47d4..948bbd68601a26e91afc61af3fe771cc85ef2b00 100644
--- a/src/pystencilssfg/lang/expressions.py
+++ b/src/pystencilssfg/lang/expressions.py
@@ -132,6 +132,9 @@ class AugExpr:
else:
return str(self._bound)
+ def __repr__(self) -> str:
+ return str(self)
+
def _bind(self, expr: DependentExpression):
if self._bound is not None:
raise SfgException("Attempting to bind an already-bound AugExpr.")