Singledispatch entrypoints

flopy4/singledispatch/plot.py

@@ -0,0 +1,9 @@
+from functools import singledispatch
+from typing import Any
+
+
+@singledispatch
+def plot(obj, **kwargs) -> Any:
+    raise NotImplementedError(
+        "plot method not implemented for type {}".format(type(obj))
+    )

flopy4/singledispatch/plot_int.py

@@ -0,0 +1,9 @@
+from typing import Any
+
+from flopy4.singledispatch.plot import plot
+
+
+@plot.register
+def _(v: int, **kwargs) -> Any:
+    print(f"Plotting a model with kwargs: {kwargs}")
+    return v

pyproject.toml

@@ -112,6 +112,9 @@ ignore = [
     "E741", # ambiguous variable name
 ]
 
+[project.entry-points.flopy4]
+plot = "flopy4.singledispatch.plot_int"
+
 [tool.pixi.project]
 channels = ["conda-forge"]
 platforms = ["win-64", "linux-64", "osx-64"]
@@ -145,3 +148,4 @@ test = { cmd = "pytest -v -n auto" }
 
 [tool.pixi.feature.lint.tasks]
 lint = { cmd = "ruff check ." }
+

test/test_singledispatch.py

@@ -0,0 +1,62 @@
+import ast
+import inspect
+import subprocess
+import sys
+from importlib.metadata import entry_points
+
+import pytest
+
+from flopy4.singledispatch.plot import plot
+
+
+def get_function_body(func):
+    source = inspect.getsource(func)
+    parsed = ast.parse(source)
+    for node in ast.walk(parsed):
+        if isinstance(node, ast.FunctionDef):
+            return ast.get_source_segment(source, node.body[0])
+    raise ValueError("Function body not found")
+
+
+def run_test_in_subprocess(test_func):
+    def wrapper():
+        test_func_source = get_function_body(test_func)
+        test_code = f"""
+import pytest
+from importlib.metadata import entry_points
+from flopy4.singledispatch.plot import plot
+
+{test_func_source}
+
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", test_code], capture_output=True, text=True
+        )
+        if result.returncode != 0:
+            print(result.stdout)
+            print(result.stderr)
+        assert result.returncode == 0, f"Test failed: {test_func.__name__}"
+
+    return wrapper
+
+
+@run_test_in_subprocess
+def test_register_singledispatch_with_entrypoints():
+    eps = entry_points(group="flopy4", name="plot")
+    for ep in eps:
+        ep.load()
+
+    # should not throw an error, because plot_int was loaded via entry points
+    return_val = plot(5)
+    assert return_val == 5
+    with pytest.raises(NotImplementedError):
+        plot("five")
+
+
+@run_test_in_subprocess
+def test_register_singledispatch_without_entrypoints():
+    # should throw an error, because plot_int was not loaded via entry points
+    with pytest.raises(NotImplementedError):
+        plot(5)
+    with pytest.raises(NotImplementedError):
+        plot("five")