get env to have all the arguments before and after context is set up

1 files changed, 8 insertions, 4 deletions

diff --git a/alembic/script.py b/alembic/script.py
index 64bc226..4c4b365 100644
--- a/alembic/script.py
+++ b/alembic/script.py
@@ -49,15 +49,19 @@ class ScriptDirectory(object):
                         yield sc
 
     def _get_rev(self, id_):
-        if id_ == 'head':
-            id_ = self._current_head()
-        elif id_ == 'base':
-            id_ = None
+        id_ = self._as_rev_number(id_)
         try:
             return self._revision_map[id_]
         except KeyError:
             raise util.CommandError("No such revision %s" % id_)
 
+    def _as_rev_number(self, id_):
+        if id_ == 'head':
+            id_ = self._current_head()
+        elif id_ == 'base':
+            id_ = None
+        return id_
+
     def _revs(self, upper, lower):
         lower = self._get_rev(lower)
         upper = self._get_rev(upper)