multimethods: avoid calling caller_locals() in Python 3 (#13238)

Python 3 metaclasses have a `__prepare__` method that lets us save the
class's dictionary before it is constructed.  In Python 2 we had to walk
up the stack using our `caller_locals()` method to get at this.  Using
`__prepare__` is much faster as it doesn't require us to use `inspect`.

This makes multimethods use the faster `__prepare__` method in Python3,
while still using `caller_locals()` in Python 2.  We try to reduce the
use of caller locals using caching to speed up Python 2 a little bit.

2 files changed, 35 insertions, 10 deletions

diff --git a/lib/spack/spack/multimethod.py b/lib/spack/spack/multimethod.py
index 085f243174..e0854a01c9 100644
--- a/lib/spack/spack/multimethod.py
+++ b/lib/spack/spack/multimethod.py
@@ -24,6 +24,7 @@ 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
 
@@ -34,6 +35,24 @@ import spack.error
 from spack.spec import Spec
 
 
+class MultiMethodMeta(type):
+    """This allows us to track the class's dict during instantiation."""
+
+    #: saved dictionary of attrs on the class being constructed
+    _locals = None
+
+    @classmethod
+    def __prepare__(cls, name, bases, **kwargs):
+        """Save the dictionary that will be used for the class namespace."""
+        MultiMethodMeta._locals = dict()
+        return MultiMethodMeta._locals
+
+    def __init__(cls, name, bases, attr_dict):
+        """Clear out the cached locals dict once the class is built."""
+        MultiMethodMeta._locals = None
+        super(MultiMethodMeta, cls).__init__(name, bases, attr_dict)
+
+
 class SpecMultiMethod(object):
     """This implements a multi-method for Spack specs.  Packages are
        instantiated with a particular spec, and you may want to
@@ -149,14 +168,15 @@ class when(object):
               def install(self, prefix):
                   # Do default install
 
-              @when('arch=chaos_5_x86_64_ib')
+              @when('target=x86_64:')
               def install(self, prefix):
                   # This will be executed instead of the default install if
-                  # the package's platform() is chaos_5_x86_64_ib.
+                  # the package's target is in the x86_64 family.
 
-              @when('arch=bgqos_0")
+              @when('target=ppc64:')
               def install(self, prefix):
-                  # This will be executed if the package's sys_type is bgqos_0
+                  # This will be executed if the package's target is in
+                  # the ppc64 family
 
        This allows each package to have a default version of install() AND
        specialized versions for particular platforms.  The version that is
@@ -202,11 +222,14 @@ class when(object):
             self.spec = Spec(condition)
 
     def __call__(self, method):
-        # Get the first definition of the method in the calling scope
-        original_method = caller_locals().get(method.__name__)
-
-        # Create a multimethod out of the original method if it
-        # isn't one already.
+        # 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__.
+        if MultiMethodMeta._locals is None:
+            MultiMethodMeta._locals = caller_locals()
+
+        # Create a multimethod with this name if there is not one already
+        original_method = MultiMethodMeta._locals.get(method.__name__)
         if not type(original_method) == SpecMultiMethod:
             original_method = SpecMultiMethod(original_method)
 
diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index dbe6e38fa6..ecff265508 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -42,6 +42,7 @@ import spack.fetch_strategy as fs
 import spack.hooks
 import spack.mirror
 import spack.mixins
+import spack.multimethod
 import spack.repo
 import spack.url
 import spack.util.web
@@ -138,7 +139,8 @@ class InstallPhase(object):
 
 class PackageMeta(
     spack.directives.DirectiveMeta,
-    spack.mixins.PackageMixinsMeta
+    spack.mixins.PackageMixinsMeta,
+    spack.multimethod.MultiMethodMeta
 ):
     """
     Package metaclass for supporting directives (e.g., depends_on) and phases