From 54617d0421f8a4c5712d548c0be46406d361c455 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9rome=20Eertmans?= <jeertmans@icloud.com>
Date: Fri, 12 Jan 2024 16:44:01 +0100
Subject: [PATCH] chore(test): add doctest namespace (#25)

* chore(test): add doctest namespace

Import JAX and NumPy in the doctest namespace to make examples less crowded by imports.

* chore(ci): apply lint suggestions

* fix(deps): conftest is now correctly removed

* fix(test): impossible to test conftest not present

* chore(ci): do not cover conftest

* chore(lint): apply fixes

* fix(ci): typo

* fix(ci): omit conftest

* chore(docs): remove trailing imports
---
 docs/source/notebooks/plotting_backend.ipynb |  2 +-
 docs/source/notebooks/quickstart.ipynb       |  2 +-
 pyproject.toml                               | 14 +++++++++++++-
 python/differt/conftest.py                   | 12 ++++++++++++
 python/differt/geometry/utils.py             |  2 +-
 python/differt/rt/image_method.py            |  3 ---
 python/differt/utils.py                      |  2 --
 7 files changed, 28 insertions(+), 9 deletions(-)
 create mode 100644 python/differt/conftest.py

diff --git a/docs/source/notebooks/plotting_backend.ipynb b/docs/source/notebooks/plotting_backend.ipynb
index 63448c21..e91b6931 100644
--- a/docs/source/notebooks/plotting_backend.ipynb
+++ b/docs/source/notebooks/plotting_backend.ipynb
@@ -182,7 +182,7 @@
    },
    "outputs": [],
    "source": [
-    "import differt.plotting as dplt\n",
+    "import differt.plotting as dplt  # noqa: E402\n",
     "\n",
     "dplt.use(\"plotly\")\n",
     "\n",
diff --git a/docs/source/notebooks/quickstart.ipynb b/docs/source/notebooks/quickstart.ipynb
index 7ed93f3d..3778e1fa 100644
--- a/docs/source/notebooks/quickstart.ipynb
+++ b/docs/source/notebooks/quickstart.ipynb
@@ -17,8 +17,8 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "import numpy as np\n",
     "import jax.numpy as jnp\n",
+    "import numpy as np\n",
     "\n",
     "import differt.plotting as dplt\n",
     "from differt.geometry import TriangleMesh\n",
diff --git a/pyproject.toml b/pyproject.toml
index a96db846..5ac984dc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,6 @@ dependencies = [
 ]
 description = "Differentiable Ray Tracing Toolbox for Radio Propagation Simulations"
 dynamic = ["license", "readme", "version"]
-exclude = ["docs/", "static/"]
 keywords = ["ray tracing", "differentiable", "propagation", "radio", "jax"]
 name = "DiffeRT"
 requires-python = ">=3.9"
@@ -58,8 +57,19 @@ exclude_lines = [
 ]
 precision = 2
 
+[tool.coverage.run]
+omit = ["**/*/conftest.py"]
+
 [tool.maturin]
 bindings = "pyo3"
+exclude = [
+  ".*",
+  ".github/**/*",
+  "benches/**/*",
+  "conftest.py",  # Path relative to source (python/differt/conftest.py does not work)
+  "docs/**/*",
+  "tests/**/*",
+]
 features = ["pyo3/extension-module"]
 module-name = "differt._core"
 python-source = "python"
@@ -138,7 +148,9 @@ target-version = "py39"
 extend-ignore-names = ["T"]
 
 [tool.ruff.lint.per-file-ignores]
+"**.ipynb" = ["B018"]
 "**/{tests,docs}/*" = ["D"]
+"python/differt/conftest.py" = ["D"]
 
 [tool.ruff.pydocstyle]
 convention = "google"
diff --git a/python/differt/conftest.py b/python/differt/conftest.py
new file mode 100644
index 00000000..30552c24
--- /dev/null
+++ b/python/differt/conftest.py
@@ -0,0 +1,12 @@
+from typing import Any
+
+import jax
+import numpy
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def add_doctest_modules(doctest_namespace: dict[str, Any]) -> None:
+    doctest_namespace["jax"] = jax
+    doctest_namespace["jnp"] = jax.numpy
+    doctest_namespace["np"] = numpy
diff --git a/python/differt/geometry/utils.py b/python/differt/geometry/utils.py
index fe80797c..9a25aa2a 100644
--- a/python/differt/geometry/utils.py
+++ b/python/differt/geometry/utils.py
@@ -37,7 +37,7 @@ def normalize(
     :Examples:
 
     >>> from differt.geometry.utils import normalize
-    >>> import jax.numpy as jnp
+    >>>
     >>> vector = jnp.array([1.0, 1.0, 1.0])
     >>> normalize(vector)  # [1., 1., 1.] / sqrt(3), sqrt(3)
     (Array([0.57735026, 0.57735026, 0.57735026], dtype=float32),
diff --git a/python/differt/rt/image_method.py b/python/differt/rt/image_method.py
index 812ac4e5..34b6f6c9 100644
--- a/python/differt/rt/image_method.py
+++ b/python/differt/rt/image_method.py
@@ -45,7 +45,6 @@ def image_of_vertices_with_respect_to_mirrors(
         a batch of random vertices. Here, normal vectors do not have a unit length,
         but they should have if you want an interpretable result.
 
-        >>> import jax
         >>> from differt.rt.image_method import (
         ...     image_of_vertices_with_respect_to_mirrors,
         ... )
@@ -137,8 +136,6 @@ def image_method(
 
         .. code-block:: python
 
-            import jax.numpy as jnp
-
             paths = image_method(
                 from_vertices, to_vertices, mirror_vertices, mirror_normals
             )
diff --git a/python/differt/utils.py b/python/differt/utils.py
index adc26adf..9c4d3bec 100644
--- a/python/differt/utils.py
+++ b/python/differt/utils.py
@@ -18,8 +18,6 @@ def sorted_array2(array: Shaped[Array, "m n"]) -> Shaped[Array, "m n"]:
     Examples:
         The following example shows how the sorting works.
 
-        >>> import jax
-        >>> import jax.numpy as jnp
         >>> from differt.utils import sorted_array2
         >>>
         >>> arr = jnp.arange(10).reshape(5, 2)