From ab11391606103ad52cf3684cc8e7ee26d52d9a5f Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Thu, 17 Oct 2024 14:56:27 +0200
Subject: [PATCH] Extend doctests
---
conftest.py | 11 ++++
docs/source/conf.py | 8 +++
pytest.ini | 8 +++
src/pystencilssfg/composer/basic_composer.py | 55 +++++++++++++++++---
src/pystencilssfg/ir/source_components.py | 2 +-
src/pystencilssfg/lang/expressions.py | 3 ++
6 files changed, 80 insertions(+), 7 deletions(-)
create mode 100644 conftest.py
create mode 100644 pytest.ini
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000..e1d0cdd
--- /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 84c2b77..11d64fd 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 0000000..94a3a6c
--- /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 7fda6c9..6a24720 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 2eab993..859f926 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 ca5a78c..948bbd6 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.")
--
GitLab