Add `when` context manager to group common constraints in packages (#24650)

This PR adds a context manager that permit to group the common part of a `when=` argument and add that to the context:
```python
class Gcc(AutotoolsPackage):
    with when('+nvptx'):
        depends_on('cuda')
        conflicts('@:6', msg='NVPTX only supported in gcc 7 and above')
        conflicts('languages=ada')
        conflicts('languages=brig')
        conflicts('languages=go')
```
The above snippet is equivalent to:
```python
class Gcc(AutotoolsPackage):
    depends_on('cuda', when='+nvptx')
    conflicts('@:6', when='+nvptx', msg='NVPTX only supported in gcc 7 and above')
    conflicts('languages=ada', when='+nvptx')
    conflicts('languages=brig', when='+nvptx')
    conflicts('languages=go', when='+nvptx')
```
which needs a repetition of the `when='+nvptx'` argument. The context manager might help improving readability and permits to group together directives related to the same semantic aspect (e.g. all the directives needed to model the behavior of `gcc` when `+nvptx` is active). 

Modifications:

- [x] Added a `when` context manager to be used with package directives
- [x] Add unit tests and documentation for the new feature
- [x] Modified `cp2k` and `gcc` to show the use of the context manager 

4 files changed, 183 insertions, 63 deletions

diff --git a/lib/spack/docs/packaging_guide.rst b/lib/spack/docs/packaging_guide.rst
index 80ddf98c38..59dad4f966 100644
--- a/lib/spack/docs/packaging_guide.rst
+++ b/lib/spack/docs/packaging_guide.rst
@@ -1257,7 +1257,7 @@ Variants
 Many software packages can be configured to enable optional
 features, which often come at the expense of additional dependencies or
 longer build times. To be flexible enough and support a wide variety of
-use cases, Spack permits to expose to the end-user the ability to choose
+use cases, Spack allows you to expose to the end-user the ability to choose
 which features should be activated in a package at the time it is installed.
 The mechanism to be employed is the :py:func:`spack.directives.variant` directive.
 
@@ -2775,6 +2775,57 @@ packages be built with MVAPICH and GCC.
 
 See the :ref:`concretization-preferences` section for more details.
 
+
+.. _group_when_spec:
+
+----------------------------
+Common ``when=`` constraints
+----------------------------
+
+In case a package needs many directives to share the whole ``when=``
+argument, or just part of it, Spack allows you to group the common part
+under a context manager:
+
+.. code-block:: python
+
+   class Gcc(AutotoolsPackage):
+
+       with when('+nvptx'):
+           depends_on('cuda')
+           conflicts('@:6', msg='NVPTX only supported in gcc 7 and above')
+           conflicts('languages=ada')
+           conflicts('languages=brig')
+           conflicts('languages=go')
+
+The snippet above is equivalent to the more verbose:
+
+.. code-block:: python
+
+   class Gcc(AutotoolsPackage):
+
+       depends_on('cuda', when='+nvptx')
+       conflicts('@:6', when='+nvptx', msg='NVPTX only supported in gcc 7 and above')
+       conflicts('languages=ada', when='+nvptx')
+       conflicts('languages=brig', when='+nvptx')
+       conflicts('languages=go', when='+nvptx')
+
+Constraints stemming from the context are added to what is explicitly present in the
+``when=`` argument of a directive, so:
+
+.. code-block:: python
+
+   with when('+elpa'):
+       depends_on('elpa+openmp', when='+openmp')
+
+is equivalent to:
+
+.. code-block:: python
+
+   depends_on('elpa+openmp', when='+openmp+elpa')
+
+Constraints from nested context managers are also added together, but they are rarely
+needed or recommended.
+
 .. _install-method:
 
 ------------------
diff --git a/lib/spack/spack/directives.py b/lib/spack/spack/directives.py
index 30a45f1b06..f9cdc261d3 100644
--- a/lib/spack/spack/directives.py
+++ b/lib/spack/spack/directives.py
@@ -27,24 +27,22 @@ The available directives are:
   * ``version``
 
 """
-
 import functools
 import os.path
 import re
 import sys
+from typing import List, Set  # novm
 
-from six import string_types
-from typing import Set, List  # novm
+import six
 
 import llnl.util.lang
 import llnl.util.tty.color
-
 import spack.error
 import spack.patch
 import spack.spec
 import spack.url
 import spack.variant
-from spack.dependency import Dependency, default_deptype, canonical_deptype
+from spack.dependency import Dependency, canonical_deptype, default_deptype
 from spack.fetch_strategy import from_kwargs
 from spack.resource import Resource
 from spack.version import Version, VersionChecksumError
@@ -114,6 +112,7 @@ class DirectiveMeta(type):
     # Set of all known directives
     _directive_names = set()  # type: Set[str]
     _directives_to_be_executed = []  # type: List[str]
+    _when_constraints_from_context = []  # type: List[str]
 
     def __new__(cls, name, bases, attr_dict):
         # Initialize the attribute containing the list of directives
@@ -168,6 +167,16 @@ class DirectiveMeta(type):
         super(DirectiveMeta, cls).__init__(name, bases, attr_dict)
 
     @staticmethod
+    def push_to_context(when_spec):
+        """Add a spec to the context constraints."""
+        DirectiveMeta._when_constraints_from_context.append(when_spec)
+
+    @staticmethod
+    def pop_from_context():
+        """Pop the last constraint from the context"""
+        return DirectiveMeta._when_constraints_from_context.pop()
+
+    @staticmethod
     def directive(dicts=None):
         """Decorator for Spack directives.
 
@@ -205,15 +214,16 @@ class DirectiveMeta(type):
         This is just a modular way to add storage attributes to the
         Package class, and it's how Spack gets information from the
         packages to the core.
-
         """
         global __all__
 
-        if isinstance(dicts, string_types):
+        if isinstance(dicts, six.string_types):
             dicts = (dicts, )
+
         if not isinstance(dicts, Sequence):
             message = "dicts arg must be list, tuple, or string. Found {0}"
             raise TypeError(message.format(type(dicts)))
+
         # Add the dictionary names if not already there
         DirectiveMeta._directive_names |= set(dicts)
 
@@ -223,6 +233,23 @@ class DirectiveMeta(type):
 
             @functools.wraps(decorated_function)
             def _wrapper(*args, **kwargs):
+                # Inject when arguments from the context
+                if DirectiveMeta._when_constraints_from_context:
+                    # Check that directives not yet supporting the when= argument
+                    # are not used inside the context manager
+                    if decorated_function.__name__ in ('version', 'variant'):
+                        msg = ('directive "{0}" cannot be used within a "when"'
+                               ' context since it does not support a "when=" '
+                               'argument')
+                        msg = msg.format(decorated_function.__name__)
+                        raise DirectiveError(msg)
+
+                    when_spec_from_context = ' '.join(
+                        DirectiveMeta._when_constraints_from_context
+                    )
+                    when_spec = kwargs.get('when', '') + ' ' + when_spec_from_context
+                    kwargs['when'] = when_spec
+
                 # If any of the arguments are executors returned by a
                 # directive passed as an argument, don't execute them
                 # lazily. Instead, let the called directive handle them.
@@ -331,7 +358,7 @@ def _depends_on(pkg, spec, when=None, type=default_deptype, patches=None):
         patches = [patches]
 
     # auto-call patch() directive on any strings in patch list
-    patches = [patch(p) if isinstance(p, string_types) else p
+    patches = [patch(p) if isinstance(p, six.string_types) else p
                for p in patches]
     assert all(callable(p) for p in patches)
 
diff --git a/lib/spack/spack/multimethod.py b/lib/spack/spack/multimethod.py
index d9fe6882b9..6ffc50f5d4 100644
--- a/lib/spack/spack/multimethod.py
+++ b/lib/spack/spack/multimethod.py
@@ -24,14 +24,12 @@ avoids overly complicated rat nests of if statements.  Obviously,
 depending on the scenario, regular old conditionals might be clearer,
 so package authors should use their judgement.
 """
-
 import functools
 import inspect
 
-from llnl.util.lang import caller_locals
-
-import spack.architecture
+import spack.directives
 import spack.error
+from llnl.util.lang import caller_locals
 from spack.spec import Spec
 
 
@@ -156,72 +154,80 @@ class SpecMultiMethod(object):
 
 
 class when(object):
-    """This annotation lets packages declare multiple versions of
-       methods like install() that depend on the package's spec.
-       For example:
+    def __init__(self, condition):
+        """Can be used both as a decorator, for multimethods, or as a context
+        manager to group ``when=`` arguments together.
 
-       .. code-block:: python
+        Examples are given in the docstrings below.
 
-          class SomePackage(Package):
-              ...
+        Args:
+            condition (str): condition to be met
+        """
+        if isinstance(condition, bool):
+            self.spec = Spec() if condition else None
+        else:
+            self.spec = Spec(condition)
 
-              def install(self, prefix):
-                  # Do default install
+    def __call__(self, method):
+        """This annotation lets packages declare multiple versions of
+        methods like install() that depend on the package's spec.
 
-              @when('target=x86_64:')
-              def install(self, prefix):
-                  # This will be executed instead of the default install if
-                  # the package's target is in the x86_64 family.
+        For example:
 
-              @when('target=ppc64:')
-              def install(self, prefix):
-                  # This will be executed if the package's target is in
-                  # the ppc64 family
+           .. code-block:: python
 
-       This allows each package to have a default version of install() AND
-       specialized versions for particular platforms.  The version that is
-       called depends on the architecutre of the instantiated package.
+              class SomePackage(Package):
+                  ...
 
-       Note that this works for methods other than install, as well.  So,
-       if you only have part of the install that is platform specific, you
-       could do this:
+                  def install(self, prefix):
+                      # Do default install
 
-       .. code-block:: python
+                  @when('target=x86_64:')
+                  def install(self, prefix):
+                      # This will be executed instead of the default install if
+                      # the package's target is in the x86_64 family.
 
-          class SomePackage(Package):
-              ...
-              # virtual dependence on MPI.
-              # could resolve to mpich, mpich2, OpenMPI
-              depends_on('mpi')
+                  @when('target=ppc64:')
+                  def install(self, prefix):
+                      # This will be executed if the package's target is in
+                      # the ppc64 family
 
-              def setup(self):
-                  # do nothing in the default case
-                  pass
+           This allows each package to have a default version of install() AND
+           specialized versions for particular platforms.  The version that is
+           called depends on the architecutre of the instantiated package.
 
-              @when('^openmpi')
-              def setup(self):
-                  # do something special when this is built with OpenMPI for
-                  # its MPI implementations.
+           Note that this works for methods other than install, as well.  So,
+           if you only have part of the install that is platform specific, you
+           could do this:
 
+           .. code-block:: python
 
-              def install(self, prefix):
-                  # Do common install stuff
-                  self.setup()
-                  # Do more common install stuff
+              class SomePackage(Package):
+                  ...
+                  # virtual dependence on MPI.
+                  # could resolve to mpich, mpich2, OpenMPI
+                  depends_on('mpi')
 
-       Note that the default version of decorated methods must
-       *always* come first.  Otherwise it will override all of the
-       platform-specific versions.  There's not much we can do to get
-       around this because of the way decorators work.
-    """
+                  def setup(self):
+                      # do nothing in the default case
+                      pass
 
-    def __init__(self, condition):
-        if isinstance(condition, bool):
-            self.spec = Spec() if condition else None
-        else:
-            self.spec = Spec(condition)
+                  @when('^openmpi')
+                  def setup(self):
+                      # do something special when this is built with OpenMPI for
+                      # its MPI implementations.
 
-    def __call__(self, method):
+
+                  def install(self, prefix):
+                      # Do common install stuff
+                      self.setup()
+                      # Do more common install stuff
+
+           Note that the default version of decorated methods must
+           *always* come first.  Otherwise it will override all of the
+           platform-specific versions.  There's not much we can do to get
+           around this because of the way decorators work.
+        """
         # In Python 2, Get the first definition of the method in the
         # calling scope by looking at the caller's locals. In Python 3,
         # we handle this using MultiMethodMeta.__prepare__.
@@ -238,6 +244,32 @@ class when(object):
 
         return original_method
 
+    def __enter__(self):
+        """Inject the constraint spec into the `when=` argument of directives
+        in the context.
+
+        This context manager allows you to write:
+
+            with when('+nvptx'):
+                conflicts('@:6', msg='NVPTX only supported from gcc 7')
+                conflicts('languages=ada')
+                conflicts('languages=brig')
+
+        instead of writing:
+
+             conflicts('@:6', when='+nvptx', msg='NVPTX only supported from gcc 7')
+             conflicts('languages=ada', when='+nvptx')
+             conflicts('languages=brig', when='+nvptx')
+
+        Context managers can be nested (but this is not recommended for readability)
+        and add their constraint to whatever may be already present in the directive
+        `when=` argument.
+        """
+        spack.directives.DirectiveMeta.push_to_context(str(self.spec))
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        spack.directives.DirectiveMeta.pop_from_context()
+
 
 class MultiMethodError(spack.error.SpackError):
     """Superclass for multimethod dispatch errors"""
diff --git a/lib/spack/spack/test/directives.py b/lib/spack/spack/test/directives.py
index 55ec46d01d..a40eff6b71 100644
--- a/lib/spack/spack/test/directives.py
+++ b/lib/spack/spack/test/directives.py
@@ -32,3 +32,13 @@ def test_true_directives_exist(mock_packages):
 
     assert cls.patches
     assert Spec() in cls.patches
+
+
+def test_constraints_from_context(mock_packages):
+    pkg_cls = spack.repo.path.get_pkg_class('with-constraint-met')
+
+    assert pkg_cls.dependencies
+    assert Spec('@1.0') in pkg_cls.dependencies['b']
+
+    assert pkg_cls.conflicts
+    assert (Spec('@1.0'), None) in pkg_cls.conflicts['%gcc']