[PATCH 1 of 3 v3] demandimport: suppport rejecting modules based on requested properties

[timeless]

# HG changeset patch
# User timeless <timeless at mozdev.org>
# Date 1474484347 0
#      Wed Sep 21 18:59:07 2016 +0000
# Node ID 33884775ee4b22109085387b0317aa17db73c483
# Parent  982fe7cdb28bb263a96b1bc2c9c3b8aedb025ab6
# Available At https://bitbucket.org/timeless/mercurial-crew
#              hg pull https://bitbucket.org/timeless/mercurial-crew -r 33884775ee4b
demandimport: suppport rejecting modules based on requested properties

Some code uses:
try:
 from Foo import Bar
except ImportError:
 from Foo import Bar2

demandimport exposes Bar as an unloaded module without checking to
see whether or not it really exists this enables it to improve loading
speed.

Unfortunately, any code that expects to get an ImportError won't get
one.

This code enables callers to tell demandimport about the shape of some
libraries to enable it to replicate the exception that callers expect.

without demandimport:
>>> from contextlib import _GeneratorContextManager
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: cannot import name _GeneratorContextManager

with demandimport without reject:
>>> from __future__ import absolute_import
>>> from mercurial import demandimport
>>> demandimport.enable()
>>> from contextlib import _GeneratorContextManager
>>> _GeneratorContextManager
<unloaded module '_GeneratorContextManager'>

with demandimport and reject:
>>> from __future__ import absolute_import
>>> from mercurial import demandimport
>>> demandimport.enable()
>>> demandimport.reject("contextlib", "_GeneratorContextManager", ImportError, "cannot import name _GeneratorContextManager")
>>> from contextlib import _GeneratorContextManager
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "mercurial/demandimport.py", line 260, in _demandimport
    processfromlist(mod, name, fromlist)
  File "mercurial/demandimport.py", line 235, in processfromlist
    raise cls(msg)
ImportError: cannot import name _GeneratorContextManager

diff -r 982fe7cdb28b -r 33884775ee4b mercurial/demandimport.py
--- a/mercurial/demandimport.py	Wed Sep 21 03:39:37 2016 +0000
+++ b/mercurial/demandimport.py	Wed Sep 21 18:59:07 2016 +0000
@@ -158,6 +158,17 @@
 _pypy = '__pypy__' in sys.builtin_module_names
 
 def _demandimport(name, globals=None, locals=None, fromlist=None, level=level):
+    def checksubimports(name, fromlist):
+        # for each item in fromlist, if a module property
+        # is registered for rejection, then raise the specified exception
+        if not name in rejects:
+            return
+        reject = rejects[name]
+        for x in fromlist:
+            if x in reject:
+                cls, msg = reject[x]
+                raise cls(msg)
+
     if not locals or name in ignore or fromlist == ('*',):
         # these cases we can't really delay
         return _hgextimport(_import, name, globals, locals, fromlist, level)
@@ -165,6 +176,7 @@
         # import a [as b]
         if '.' in name: # a.b
             base, rest = name.split('.', 1)
+            checksubimports(base, [rest])
             # email.__init__ loading email.mime
             if globals and globals.get('__name__', None) == base:
                 return _import(name, globals, locals, fromlist, level)
@@ -220,6 +232,7 @@
                 mod = getattr(mod, comp)
             return mod
 
+        checksubimports(name, fromlist)
         if level >= 0:
             if name:
                 # "from a import b" or "from .a import b" style
@@ -287,6 +300,16 @@
     'distutils.msvc9compiler',
     ]
 
+rejects = {}
+
+def reject(mod, prop, cls, msg):
+    """inform demandimport that a module does not have a property
+
+    arguments are the class and message to raise when code tries to
+    import it."""
+    if not mod in rejects:
+        rejects[mod] = {}
+    rejects[mod][prop] = [cls, msg]
 def isenabled():
     return builtins.__import__ == _demandimport
 
diff -r 982fe7cdb28b -r 33884775ee4b tests/test-demandimport.py
--- a/tests/test-demandimport.py	Wed Sep 21 03:39:37 2016 +0000
+++ b/tests/test-demandimport.py	Wed Sep 21 18:59:07 2016 +0000
@@ -63,6 +63,14 @@
 print("re.stderr =", f(re.stderr))
 print("re =", f(re))
 
+demandimport.reject('chunk', 'Chunk', ImportError, 'boo')
+try:
+    from chunk import Chunk
+    print('reject should prevent chunk from loading for Chunk')
+    Chunk.getname
+except ImportError:
+    pass
+
 demandimport.disable()
 os.environ['HGDEMANDIMPORT'] = 'disable'
 # this enable call should not actually enable demandimport!