Singledispatch entrypoints (#38)

I was playing with how we could set up functions in a modular fashion without having to add a function in the corresponding class. cattrs mentions the use of single dispatch functions for their structure and unstructure functions. We can do the same for functions like plot(), where the plot function should be seen as a module, and not as part of the class itself.

I have combined this functionality with the entry_points function to show that a user or plugin package could in theory also add functionality. For example, a user comes with a specific entry point for plotting a specific package, other than the default way that we normally do.

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")