ENH add mpi virtual packages

.gitignore

@@ -25,3 +25,5 @@ coverage.xml
 .ipynb_checkpoints/
 .pytest_cache/*
 .pytest_cache
+
+conda_forge_conda_plugins/_version.py

README.md

@@ -7,4 +7,8 @@ conda plugins for the `conda-forge` ecosystem
 
 ### OpenMPI Virtual Package
 
-This plugin provides the `__openmpi` virtual package to be used in conjunction with the `external*` builds of the `conda-forge` `openmpi` package.
+This plugin provides the `__openmpi` virtual package to be used in conjunction with the `external*` builds of the `conda-forge` `openmpi` package. If the `ompi_info` command is found in the path, the `__openmpi` virtual package is set to the installed version of OpenMPI. You can override the version by setting the `CONDA_OVERRIDE_OPENMPI` environment variable.
+
+### MPICH Virtual Package
+
+This plugin provides the `__mpich` virtual package to be used in conjunction with the `external*` builds of the `conda-forge` `mpich` package. If the `mpichversion` command is found in the path, the `__mpich` virtual package is set to the installed version of MPICH. You can override the version by setting the `CONDA_OVERRIDE_MPICH` environment variable.

conda_forge_conda_plugins/hooks.py

@@ -1,11 +1,55 @@
-from conda.plugins import CondaSubcommand, hookimpl
+import os
+import subprocess
+
+from conda.plugins import CondaVirtualPackage, hookimpl
 
 
 @hookimpl
-def conda_subcommands():
-    def hello_conda(args):
-        print("Hello conda!")
+def conda_virtual_packages():
+    # openmpi virtual package
+    openmpi_version = None
+    if "CONDA_OVERRIDE_OPENMPI" in os.environ:
+        openmpi_version = os.environ["CONDA_OVERRIDE_OPENMPI"]
+    else:
+        try:
+            ret = subprocess.run(
+                ["ompi_info", "--parsable"], capture_output=True, text=True
+            )
+            if ret.returncode == 0:
+                for line in ret.stdout.splitlines():
+                    if line.startswith("ompi:version:full:"):
+                        openmpi_version = line.strip().split(":")[3].strip()
+                        break
+        except Exception:
+            pass
+
+    if openmpi_version:
+        yield CondaVirtualPackage(
+            name="openmpi",
+            version=openmpi_version,
+            build=None,
+        )
+
+    # mpich virtual package
+    mpich_version = None
+    if "CONDA_OVERRIDE_MPICH" in os.environ:
+        mpich_version = os.environ["CONDA_OVERRIDE_MPICH"]
+    else:
+        try:
+            ret = subprocess.run(
+                ["mpichversion", "--version"], capture_output=True, text=True
+            )
+            if ret.returncode == 0:
+                for line in ret.stdout.splitlines():
+                    if line.startswith("MPICH Version:"):
+                        mpich_version = line.strip().split(":")[1].strip()
+                        break
+        except Exception:
+            pass
 
-    yield CondaSubcommand(
-        name="hello", action=hello_conda, summary='Command that prints "Hello conda!"'
-    )
+    if mpich_version:
+        yield CondaVirtualPackage(
+            name="mpich",
+            version=mpich_version,
+            build=None,
+        )

tests/test_mpich.py

@@ -0,0 +1,23 @@
+import json
+import os
+import subprocess
+
+
+def test_mpich_override():
+    env = {"CONDA_OVERRIDE_MPICH": "10.0.0"}
+    env.update(dict(os.environ))
+    ret = subprocess.run(
+        ["conda", "info", "--json"],
+        capture_output=True,
+        text=True,
+        env=env,
+    )
+    assert ret.returncode == 0
+    info = json.loads(ret.stdout)
+    pkg = None
+    for _pkg in info["virtual_pkgs"]:
+        if _pkg[0] == "__mpich":
+            pkg = _pkg
+            break
+    assert pkg is not None, info["virtual_pkgs"]
+    assert pkg[1] == "10.0.0", pkg

tests/test_openmpi.py

@@ -0,0 +1,23 @@
+import json
+import os
+import subprocess
+
+
+def test_openmpi_override():
+    env = {"CONDA_OVERRIDE_OPENMPI": "10.0.0"}
+    env.update(dict(os.environ))
+    ret = subprocess.run(
+        ["conda", "info", "--json"],
+        capture_output=True,
+        text=True,
+        env=env,
+    )
+    assert ret.returncode == 0
+    info = json.loads(ret.stdout)
+    pkg = None
+    for _pkg in info["virtual_pkgs"]:
+        if _pkg[0] == "__openmpi":
+            pkg = _pkg
+            break
+    assert pkg is not None, info["virtual_pkgs"]
+    assert pkg[1] == "10.0.0", pkg