[PATCH] demandimport state stack

[yozh at mx1.ru]

# HG changeset patch
# User Stepan Koltsov <stepancheg at yandex-team.ru>
# Date 1241277646 -14400
# Node ID ac16f3645ca9fd9745d026e8d6435c6e5b717b92
# Parent  e3d3dad805f9f7c5f17c7a80cd659ecf25238ca1
allow stacking of demandimport state

If some extension modules wishes do disable demandimport while importing
it could call demandimport.disable() and at the end it should call demandimport.enable().
This is wrong, because demandimport could be already disabled at the start of module initialization.

So the proper way to disable demandimport from inside module is:

from mercurial import demandimport
import __builtin__
orig_import = __builtin__.__import__
try:
    demandimport.disable()
    ... regular imports
finally:
    __builtin__.__import__ = orig_import

This patch hides __builtin__ manipulation in demandimport module:

from mercurial import demandimport
demandimport.push_state()
try:
    demandimport.disable()
    ... regular imports
finally:
    demandimport.pop_state()

diff -r e3d3dad805f9 -r ac16f3645ca9 mercurial/demandimport.py
--- a/mercurial/demandimport.py	Fri May 01 11:32:19 2009 +0200
+++ b/mercurial/demandimport.py	Sat May 02 19:20:46 2009 +0400
@@ -124,6 +124,7 @@
     # imported by profile, itself imported by hotshot.stats,
     # not available under Windows
     'resource',
+    'state',
     ]
 
 def enable():
@@ -134,3 +135,28 @@
     "disable global demand-loading of modules"
     __builtin__.__import__ = _origimport
 
+state = []
+
+def push_state():
+    """
+    Push current import function onto the stack, return stack state object.
+    State can be reset by pop_state() or restore_state(state) methods.
+    """
+    global state
+    state = state + [__builtin__.__import__]
+    return state
+
+def pop_state():
+    """
+    Pop top-level import function from the internal save-state stack.
+    """
+    restore_state(state)
+
+def restore_state(old_state):
+    """
+    Restore state stack to the passed save-state object.
+    """
+    global state
+    __builtin__.__import__ = old_state[-1]
+    state = old_state[:-1]
+
diff -r e3d3dad805f9 -r ac16f3645ca9 tests/test-demandimport-stack
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-demandimport-stack	Sat May 02 19:20:46 2009 +0400
@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+#
+# Test demandimport stack manipulation.
+#
+# copy-paste and rewrite of test-demandimport
+#
+
+from mercurial import demandimport
+
+di_state = demandimport.push_state()
+
+import re
+
+rsub = re.sub
+def f(obj):
+    l = repr(obj)
+    l = rsub("0x[0-9a-fA-F]+", "0x?", l)
+    l = rsub("from '.*'", "from '?'", l)
+    return l
+
+demandimport.enable()
+
+import os
+
+print "os =", f(os)
+print "os.system =", f(os.system)
+print "os =", f(os)
+
+demandimport.push_state()
+
+demandimport.disable()
+
+from mercurial import util
+
+print "util =", f(util)
+print "util.system =", f(util.system)
+print "util =", f(util)
+print "util.system =", f(util.system)
+
+demandimport.pop_state()
+
+demandimport.enable()
+
+import re as fred
+print "fred =", f(fred)
+
+demandimport.restore_state(di_state)
+
+import sys as re
+print "re =", f(re)
+
+print "fred =", f(fred)
+print "fred.sub =", f(fred.sub)
+print "fred =", f(fred)
+
+print "re =", f(re)
+print "re.stdout =", f(re.stdout)
+print "re =", f(re)
diff -r e3d3dad805f9 -r ac16f3645ca9 tests/test-demandimport-stack.out
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-demandimport-stack.out	Sat May 02 19:20:46 2009 +0400
@@ -0,0 +1,15 @@
+os = <unloaded module 'os'>
+os.system = <built-in function system>
+os = <module 'os' from '?'>
+util = <module 'mercurial.util' from '?'>
+util.system = <function system at 0x?>
+util = <module 'mercurial.util' from '?'>
+util.system = <function system at 0x?>
+fred = <unloaded module 're'>
+re = <module 'sys' (built-in)>
+fred = <unloaded module 're'>
+fred.sub = <function sub at 0x?>
+fred = <proxied module 're'>
+re = <module 'sys' (built-in)>
+re.stdout = <open file '<stdout>', mode 'w' at 0x?>
+re = <module 'sys' (built-in)>