Add timezone option to config

Using dateutil.tz to link string names to tzinfo objects,
the create_date can now generate using a named timezone
rather than datetime.now().

Change-Id: I9f151cb9e11da3d68be63d7141f60e7eccb9812c
Fixes: #425

alembic/script/base.py

@@ -1,4 +1,5 @@
 import datetime
+from dateutil import tz
 import os
 import re
 import shutil
@@ -42,7 +43,8 @@ class ScriptDirectory(object):
     def __init__(self, dir, file_template=_default_file_template,
                  truncate_slug_length=40,
                  version_locations=None,
-                 sourceless=False, output_encoding="utf-8"):
+                 sourceless=False, output_encoding="utf-8",
+                 timezone=None):
         self.dir = dir
         self.file_template = file_template
         self.version_locations = version_locations
@@ -50,6 +52,7 @@ class ScriptDirectory(object):
         self.sourceless = sourceless
         self.output_encoding = output_encoding
         self.revision_map = revision.RevisionMap(self._load_revisions)
+        self.timezone = timezone
 
         if not os.access(dir, os.F_OK):
             raise util.CommandError("Path doesn't exist: %r.  Please use "
@@ -118,6 +121,7 @@ class ScriptDirectory(object):
         version_locations = config.get_main_option("version_locations")
         if version_locations:
             version_locations = _split_on_space_comma.split(version_locations)
+
         return ScriptDirectory(
             util.coerce_resource_to_filename(script_location),
             file_template=config.get_main_option(
@@ -126,7 +130,8 @@ class ScriptDirectory(object):
             truncate_slug_length=truncate_slug_length,
             sourceless=config.get_main_option("sourceless") == "true",
             output_encoding=config.get_main_option("output_encoding", "utf-8"),
-            version_locations=version_locations
+            version_locations=version_locations,
+            timezone=config.get_main_option("timezone")
         )
 
     @contextmanager
@@ -440,6 +445,18 @@ class ScriptDirectory(object):
                 "Creating directory %s" % path,
                 os.makedirs, path)
 
+    def _generate_create_date(self):
+        if self.timezone is not None:
+            tzinfo = tz.gettz(self.timezone.upper())
+            if tzinfo is None:
+                raise util.CommandError(
+                    "Can't locate timezone: %s" % self.timezone)
+            create_date = datetime.datetime.utcnow().replace(
+                tzinfo=tz.tzutc()).astimezone(tzinfo)
+        else:
+            create_date = datetime.datetime.now()
+        return create_date
+
     def generate_revision(
             self, revid, message, head=None,
             refresh=False, splice=False, branch_labels=None,
@@ -478,7 +495,7 @@ class ScriptDirectory(object):
         if len(set(heads)) != len(heads):
             raise util.CommandError("Duplicate head revisions specified")
 
-        create_date = datetime.datetime.now()
+        create_date = self._generate_create_date()
 
         if version_path is None:
             if len(self._version_locations) > 1:

alembic/templates/generic/alembic.ini.mako

@@ -7,6 +7,12 @@ script_location = ${script_location}
 # template used to generate migration files
 # file_template = %%(rev)s_%%(slug)s
 
+# timezone to use when rendering the date
+# within the migration file as well as the filename.
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
 # max length of characters to apply to the
 # "slug" field
 #truncate_slug_length = 40

alembic/templates/multidb/alembic.ini.mako

@@ -7,6 +7,12 @@ script_location = ${script_location}
 # template used to generate migration files
 # file_template = %%(rev)s_%%(slug)s
 
+# timezone to use when rendering the date
+# within the migration file as well as the filename.
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
 # max length of characters to apply to the
 # "slug" field
 #truncate_slug_length = 40

alembic/templates/pylons/alembic.ini.mako

@@ -7,6 +7,12 @@ script_location = ${script_location}
 # template used to generate migration files
 # file_template = %%(rev)s_%%(slug)s
 
+# timezone to use when rendering the date
+# within the migration file as well as the filename.
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
 # max length of characters to apply to the
 # "slug" field
 #truncate_slug_length = 40

docs/build/changelog.rst

@@ -7,6 +7,16 @@ Changelog
     :version: 0.9.2
     :released:
 
+    .. change:: 425
+      :tags: feature, commands
+      :tickets: 425
+
+      Added a new configuration option ``timezone``, a string timezone name
+      that will be applied to the create date timestamp rendered
+      inside the revision file as made availble to the ``file_template`` used
+      to generate the revision filename.  Note this change adds the
+      ``python-dateutil`` package as a dependency.
+
     .. change:: 421
       :tags: bug, autogenerate
       :tickets: 421

docs/build/tutorial.rst

@@ -119,6 +119,12 @@ The file generated with the "generic" configuration looks like::
     # template used to generate migration files
     # file_template = %%(rev)s_%%(slug)s
 
+    # timezone to use when rendering the date
+    # within the migration file as well as the filename.
+    # string value is passed to dateutil.tz.gettz()
+    # leave blank for localtime
+    # timezone =
+
     # max length of characters to apply to the
     # "slug" field
     #truncate_slug_length = 40
@@ -212,8 +218,23 @@ This file contains the following features:
     * ``%%(rev)s`` - revision id
     * ``%%(slug)s`` - a truncated string derived from the revision message
     * ``%%(year)d``, ``%%(month).2d``, ``%%(day).2d``, ``%%(hour).2d``,
-      ``%%(minute).2d``, ``%%(second).2d`` - components of the create date
+      ``%%(minute).2d``, ``%%(second).2d`` - components of the create date,
-      as returned by ``datetime.datetime.now()``
+      by default ``datetime.datetime.now()`` unless the ``timezone``
+      configuration option is also used.
+
+* ``timezone`` - an optional timezone name (e.g. ``UTC``, ``EST5EDT``, etc.)
+  that will be applied to the timestamp which renders inside the migration
+  file's comment as well as within the filename.  If ``timezone`` is specified,
+  the create date object is no longer derived from ``datetime.datetime.now()``
+  and is instead generated as::
+
+      datetime.datetime.utcnow().replace(
+            tzinfo=dateutil.tz.tzutc()
+      ).astimezone(
+          dateutil.tz.gettz(<timezone>)
+      )
+
+  .. versionadded:: 0.9.2
 
 * ``truncate_slug_length`` - defaults to 40, the max number of characters
   to include in the "slug" field.

setup.py

@@ -16,6 +16,7 @@ requires = [
     'SQLAlchemy>=0.7.6',
     'Mako',
     'python-editor>=0.3',
+    'python-dateutil'
 ]
 
 try:

tests/test_script_production.py

@@ -18,6 +18,7 @@ import sqlalchemy as sa
 from sqlalchemy.engine.reflection import Inspector
 from alembic.util import CommandError
 import re
+from dateutil import tz
 
 env, abc, def_ = None, None, None
 
@@ -153,6 +154,72 @@ class ScriptNamingTest(TestBase):
                 "message_2012_7_25_15_8_5.py" % _get_staging_directory())
         )
 
+    def _test_tz(self, timezone_arg, given, expected):
+        script = ScriptDirectory(
+            _get_staging_directory(),
+            file_template="%(rev)s_%(slug)s_"
+            "%(year)s_%(month)s_"
+            "%(day)s_%(hour)s_"
+            "%(minute)s_%(second)s",
+            timezone=timezone_arg
+        )
+
+        with mock.patch(
+                "alembic.script.base.datetime",
+                mock.Mock(
+                    datetime=mock.Mock(
+                        utcnow=lambda: given,
+                        now=lambda: given
+                    )
+                )
+        ):
+            create_date = script._generate_create_date()
+        eq_(
+            create_date,
+            expected
+        )
+
+    def test_custom_tz(self):
+        self._test_tz(
+            'EST5EDT',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT'))
+        )
+
+    def test_custom_tz_lowercase(self):
+        self._test_tz(
+            'est5edt',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT'))
+        )
+
+    def test_custom_tz_utc(self):
+        self._test_tz(
+            'utc',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 15, 8, 5, tzinfo=tz.gettz('UTC'))
+        )
+
+    def test_default_tz(self):
+        self._test_tz(
+            None,
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(2012, 7, 25, 15, 8, 5)
+        )
+
+    def test_tz_cant_locate(self):
+        assert_raises_message(
+            CommandError,
+            "Can't locate timezone: fake",
+            self._test_tz,
+            "fake",
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(2012, 7, 25, 15, 8, 5)
+        )
+
 
 class RevisionCommandTest(TestBase):
     def setUp(self):