turn alembic.op and alembic.context into real proxy modules,

with an accurate system of reflecting the Operations and
EnvironmentContext methods into them.

6 files changed, 93 insertions, 42 deletions

diff --git a/alembic/__init__.py b/alembic/__init__.py
index 09d91df..dd0591e 100644
--- a/alembic/__init__.py
+++ b/alembic/__init__.py
@@ -6,10 +6,5 @@ package_dir = path.abspath(path.dirname(__file__))
 
 
 from alembic import op
-
-class _ContextProxy(object):
-    """A proxy object for the current :class:`.EnvironmentContext`."""
-    def __getattr__(self, key):
-        return getattr(_context, key)
-context = _ContextProxy()
+from alembic import context
 
diff --git a/alembic/context.py b/alembic/context.py
new file mode 100644
index 0000000..6c33d85
--- /dev/null
+++ b/alembic/context.py
@@ -0,0 +1,6 @@
+from alembic.environment import EnvironmentContext
+from alembic import util
+
+# create proxy functions for 
+# each method on the EnvironmentContext class.
+util.create_module_class_proxy(EnvironmentContext, globals(), locals())
diff --git a/alembic/environment.py b/alembic/environment.py
index f61c9c7..5e86fcd 100644
--- a/alembic/environment.py
+++ b/alembic/environment.py
@@ -46,12 +46,12 @@ class EnvironmentContext(object):
         be made available as ``from alembic import context``.
     
         """
-        alembic._context = self
+        alembic.context._install_proxy(self)
         return self
 
     def __exit__(self, *arg, **kw):
-        alembic._context = None
-        alembic.op._proxy = None
+        alembic.context._remove_proxy()
+        alembic.op._remove_proxy()
 
     def is_offline_mode(self):
         """Return True if the current migrations environment 
@@ -78,7 +78,7 @@ class EnvironmentContext(object):
         made available via :meth:`.configure`.
 
         """
-        return self.migration_context.impl.transactional_ddl
+        return self.get_context().impl.transactional_ddl
 
     def requires_connection(self):
         return not self.is_offline_mode()
@@ -105,7 +105,7 @@ class EnvironmentContext(object):
 
         """
         if self._migration_context is not None:
-            return self.script._as_rev_number(self.migration_context._start_from_rev)
+            return self.script._as_rev_number(self.get_context()._start_from_rev)
         elif 'starting_rev' in self.context_opts:
             return self.script._as_rev_number(self.context_opts['starting_rev'])
         else:
@@ -345,7 +345,7 @@ class EnvironmentContext(object):
 
         """
         with Operations.context(self._migration_context):
-            self.migration_context.run_migrations(**kw)
+            self.get_context().run_migrations(**kw)
 
     def execute(self, sql):
         """Execute the given SQL using the current change context.
@@ -359,7 +359,7 @@ class EnvironmentContext(object):
         made available via :meth:`.configure`.
 
         """
-        self.migration_context.execute(sql)
+        self.get_context().execute(sql)
 
     def static_output(self, text):
         """Emit text directly to the "offline" SQL stream.
@@ -370,7 +370,7 @@ class EnvironmentContext(object):
         is added, etc.
     
         """
-        self.migration_context.impl.static_output(text)
+        self.get_context().impl.static_output(text)
 
     def begin_transaction(self):
         """Return a context manager that will 
@@ -423,30 +423,25 @@ class EnvironmentContext(object):
         elif self.is_offline_mode():
             @contextmanager
             def begin_commit():
-                self.migration_context.impl.emit_begin()
+                self.get_context().impl.emit_begin()
                 yield
-                self.migration_context.impl.emit_commit()
+                self.get_context().impl.emit_commit()
             return begin_commit()
         else:
             return self.get_bind().begin()
 
-    @property
-    def migration_context(self):
+    def get_context(self):
         """Return the current :class:`.MigrationContext` object.
 
         If :meth:`.EnvironmentContext.configure` has not been called yet, raises
         an exception.
 
         """
+
         if self._migration_context is None:
             raise Exception("No context has been configured yet.")
         return self._migration_context
 
-    def get_context(self):
-        """A synonym for :attr:`.EnvironmentContext.migration_context`."""
-
-        return self.migration_context
-
     def get_bind(self):
         """Return the current 'bind'.
 
@@ -458,9 +453,9 @@ class EnvironmentContext(object):
         made available via :meth:`.configure`.
 
         """
-        return self.migration_context.bind
+        return self.get_context().bind
 
     def get_impl(self):
-        return self.migration_context.impl
+        return self.get_context().impl
 
 configure = EnvironmentContext
diff --git a/alembic/op.py b/alembic/op.py
index 8a5e0fa..9f2a26b 100644
--- a/alembic/op.py
+++ b/alembic/op.py
@@ -1,19 +1,6 @@
 from alembic.operations import Operations
+from alembic import util
 
 # create proxy functions for 
 # each method on the Operations class.
-
-# TODO: this is a quick and dirty version of this.
-# Ideally, we'd be duplicating method signatures 
-# and such, using eval(), etc.
-
-_proxy = None
-def _create_op_proxy(name):
-    def go(*arg, **kw):
-        return getattr(_proxy, name)(*arg, **kw)
-    go.__name__ = name
-    return go
-
-for methname in dir(Operations):
-    if not methname.startswith('_'):
-        locals()[methname] = _create_op_proxy(methname)
\ No newline at end of file
+util.create_module_class_proxy(Operations, globals(), locals())
diff --git a/alembic/operations.py b/alembic/operations.py
index f3e6708..9efa830 100644
--- a/alembic/operations.py
+++ b/alembic/operations.py
@@ -37,9 +37,9 @@ class Operations(object):
     @contextmanager
     def context(cls, migration_context):
         op = Operations(migration_context)
-        alembic.op._proxy = op
+        alembic.op._install_proxy(op)
         yield op
-        del alembic.op._proxy
+        alembic.op._remove_proxy()
 
     def _foreign_key_constraint(self, name, source, referent, local_cols, remote_cols):
         m = schema.MetaData()
diff --git a/alembic/util.py b/alembic/util.py
index f58992a..3ae15f9 100644
--- a/alembic/util.py
+++ b/alembic/util.py
@@ -5,9 +5,11 @@ import sys
 import os
 import textwrap
 from sqlalchemy.engine import url
+from sqlalchemy import util as sqla_util
 import imp
 import warnings
 import re
+import inspect
 import time
 import random
 import uuid
@@ -42,6 +44,72 @@ def template_to_file(template_file, dest, **kw):
             Template(filename=template_file).render(**kw)
         )
 
+def create_module_class_proxy(cls, globals_, locals_):
+    """Create module level proxy functions for the
+    methods on a given class.
+    
+    The functions will have a compatible signature
+    as the methods.   A proxy is established
+    using the ``_install_proxy(obj)`` function,
+    and removed using ``_remove_proxy()``, both
+    installed by calling this function.
+
+    """
+    attr_names = set()
+
+    def _install_proxy(obj):
+        globals_['_proxy'] = obj
+        for name in attr_names:
+            globals_[name] = getattr(obj, name)
+
+    def _remove_proxy():
+        globals_['_proxy'] = None
+        for name in attr_names:
+            del globals_[name]
+
+    globals_['_install_proxy'] = _install_proxy
+    globals_['_remove_proxy'] = _remove_proxy
+
+    def _create_op_proxy(name):
+        fn = getattr(cls, name)
+        spec = inspect.getargspec(fn)
+        if spec[0] and spec[0][0] == 'self':
+            spec[0].pop(0)
+        args = inspect.formatargspec(*spec)
+        num_defaults = 0
+        if spec[3]:
+            num_defaults += len(spec[3])
+        name_args = spec[0]
+        if num_defaults:
+            defaulted_vals = name_args[0-num_defaults:]
+        else:
+            defaulted_vals = ()
+
+        apply_kw = inspect.formatargspec(
+                                name_args, spec[1], spec[2], 
+                                defaulted_vals,
+                                formatvalue=lambda x: '=' + x)
+
+        func_text = textwrap.dedent("""\
+        def %(name)s(%(args)s):
+            %(doc)r
+            return _proxy.%(name)s(%(apply_kw)s)
+        """ % {
+            'name':name,
+            'args':args[1:-1],
+            'apply_kw':apply_kw[1:-1],
+            'doc':fn.__doc__,
+        })
+        lcl = {}
+        exec func_text in globals_, lcl
+        return lcl[name]
+
+    for methname in dir(cls):
+        if not methname.startswith('_'):
+            if callable(getattr(cls, methname)):
+                locals_[methname] = _create_op_proxy(methname)
+            else:
+                attr_names.add(methname)
 
 def status(_statmsg, fn, *arg, **kw):
     msg(_statmsg + "...", False)