add option to customize templates and use multiple themes

9 files changed, 172 insertions, 96 deletions

diff --git a/docs/changelog.rst b/docs/changelog.rst
index f09e5c9..39acb8f 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -1,6 +1,7 @@
 0.5.5
 -----
 
+- added option to define custom templates through option ``--templates_path``, read more in :ref:`tutorial section <custom-templates>`
 - url parameter can also be an Engine instance (this usage is discouraged though sometimes necessary)
 - added support for SQLAlchemy 0.6 (missing oracle and firebird) by Michael Bayer
 - alter, create, drop column / rename table / rename index constructs now accept `alter_metadata` parameter. If True, it will modify Column/Table objects according to changes. Otherwise, everything will be untouched.
diff --git a/docs/versioning.rst b/docs/versioning.rst
index 740196f..fa5483d 100644
--- a/docs/versioning.rst
+++ b/docs/versioning.rst
@@ -493,3 +493,25 @@ currently:
   the databases your application will actually be using to ensure your
   updates to that database work properly. This must be a list;
   example: `['postgres', 'sqlite']`
+
+
+.. _custom-templates:
+
+Customize templates
+===================
+
+Users can pass ``templates_path`` to API functions to provide customized templates path.
+Path should be a collection of templates, like ``migrate.versioning.templates`` package directory.
+
+One may also want to specify custom themes. API functions accept ``templates_theme`` for this purpose (which defaults to `default`)
+
+Example::
+	
+	/home/user/templates/manage $ ls
+	default.py_tmpl
+	pylons.py_tmpl
+
+	/home/user/templates/manage $ migrate manage manage.py --templates_path=/home/user/templates --templates_theme=pylons
+
+
+.. versionadded:: 0.6.0
diff --git a/migrate/versioning/api.py b/migrate/versioning/api.py
index 6e0da11..fe45734 100644
--- a/migrate/versioning/api.py
+++ b/migrate/versioning/api.py
@@ -264,7 +264,7 @@ def manage(file, **opts):
         python manage.py version
         %prog version --repository=/path/to/repository
     """
-    return Repository.create_manage_file(file, **opts)
+    Repository.create_manage_file(file, **opts)
 
 
 def compare_model_to_db(url, model, repository, **opts):
diff --git a/migrate/versioning/repository.py b/migrate/versioning/repository.py
index 3eb2f74..0938672 100644
--- a/migrate/versioning/repository.py
+++ b/migrate/versioning/repository.py
@@ -4,10 +4,10 @@
 import os
 import shutil
 import string
-from pkg_resources import resource_string, resource_filename
+from pkg_resources import resource_filename
 
 from migrate.versioning import exceptions, script, version, pathed, cfgparse
-from migrate.versioning.template import template
+from migrate.versioning.template import Template
 from migrate.versioning.base import *
 
 
@@ -91,11 +91,18 @@ class Repository(pathed.Pathed):
         except exceptions.PathNotFoundError, e:
             raise exceptions.InvalidRepositoryError(path)
 
-    # TODO: what are those options?
     @classmethod
-    def prepare_config(cls, pkg, rsrc, name, **opts):
+    def prepare_config(cls, tmpl_dir, config_file, name, **opts):
         """
         Prepare a project configuration file for a new project.
+
+        :param tmpl_dir: Path to Repository template
+        :param config_file: Name of the config file in Repository template
+        :param name: Repository name
+        :type tmpl_dir: string
+        :type config_file: string
+        :type name: string
+        :returns: Populated config file
         """
         # Prepare opts
         defaults = dict(
@@ -105,7 +112,7 @@ class Repository(pathed.Pathed):
 
         defaults.update(opts)
 
-        tmpl = resource_string(pkg, rsrc)
+        tmpl = open(os.path.join(tmpl_dir, config_file)).read()
         ret = string.Template(tmpl).substitute(defaults)
         return ret
 
@@ -113,14 +120,12 @@ class Repository(pathed.Pathed):
     def create(cls, path, name, **opts):
         """Create a repository at a specified path"""
         cls.require_notfound(path)
-
-        pkg, rsrc = template.get_repository(as_pkg=True)
-        tmplpkg = '.'.join((pkg, rsrc))
-        tmplfile = resource_filename(pkg, rsrc)
-        config_text = cls.prepare_config(tmplpkg, cls._config, name, **opts)
+        theme = opts.get('templates_theme', None)
 
         # Create repository
-        shutil.copytree(tmplfile, path)
+        tmpl_dir = Template(opts.pop('templates_path', None)).get_repository(theme=theme)
+        config_text = cls.prepare_config(tmpl_dir, cls._config, name, **opts)
+        shutil.copytree(tmpl_dir, path)
 
         # Edit config defaults
         fd = open(os.path.join(path, cls._config), 'w')
@@ -129,7 +134,7 @@ class Repository(pathed.Pathed):
 
         # Create a management script
         manager = os.path.join(path, 'manage.py')
-        Repository.create_manage_file(manager, repository=path)
+        Repository.create_manage_file(manager, theme=theme, repository=path)
 
         return cls(path)
 
@@ -205,12 +210,10 @@ class Repository(pathed.Pathed):
         :param file_: Destination file to be written
         :param opts: Options that are passed to template
         """
+        mng_file = Template(opts.pop('templates_path', None)).get_manage(theme=opts.pop('templates_theme', None))
         vars_ = ",".join(["%s='%s'" % var for var in opts.iteritems()])
 
-        pkg, rsrc = template.manage(as_pkg=True)
-        tmpl = resource_string(pkg, rsrc)
-        result = tmpl % dict(defaults=vars_)
-
+        tmpl = open(mng_file).read()
         fd = open(file_, 'w')
-        fd.write(result)
+        fd.write(tmpl % dict(defaults=vars_))
         fd.close()
diff --git a/migrate/versioning/script/py.py b/migrate/versioning/script/py.py
index 15f68b0..b8f29a1 100644
--- a/migrate/versioning/script/py.py
+++ b/migrate/versioning/script/py.py
@@ -7,7 +7,7 @@ from StringIO import StringIO
 import migrate
 from migrate.versioning import exceptions, genmodel, schemadiff
 from migrate.versioning.base import operations
-from migrate.versioning.template import template
+from migrate.versioning.template import Template
 from migrate.versioning.script import base
 from migrate.versioning.util import import_path, load_model, construct_engine
 
@@ -22,11 +22,7 @@ class PythonScript(base.BaseScript):
         :returns: :class:`PythonScript instance <migrate.versioning.script.py.PythonScript>`"""
         cls.require_notfound(path)
 
-        # TODO: Use the default script template (defined in the template
-        # module) for now, but we might want to allow people to specify a
-        # different one later.
-        template_file = None
-        src = template.get_script(template_file)
+        src = Template(opts.pop('templates_path', None)).get_script(theme=opts.pop('templates_theme', None))
         shutil.copy(src, path)
 
         return cls(path)
@@ -67,8 +63,7 @@ class PythonScript(base.BaseScript):
             genmodel.ModelGenerator(diff).toUpgradeDowngradePython()
 
         # Store differences into file.
-        # TODO: add custom templates
-        src = template.get_script(None)
+        src = Template(opts.pop('templates_path', None)).get_script(opts.pop('templates_theme', None))
         f = open(src)
         contents = f.read()
         f.close()
diff --git a/migrate/versioning/template.py b/migrate/versioning/template.py
index 91606a7..aa1cc34 100644
--- a/migrate/versioning/template.py
+++ b/migrate/versioning/template.py
@@ -4,81 +4,84 @@
 import os
 import shutil
 import sys
+
 from pkg_resources import resource_filename
 
 from migrate.versioning.base import *
 from migrate.versioning import pathed
 
 
-class Packaged(pathed.Pathed):
-    """An object assoc'ed with a Python package"""
-
-    def __init__(self, pkg):
-        self.pkg = pkg
-        path = self._find_path(pkg)
-        super(Packaged, self).__init__(path)
-
-    @classmethod
-    def _find_path(cls, pkg):
-        pkg_name, resource_name = pkg.rsplit('.', 1)
-        ret = resource_filename(pkg_name, resource_name)
-        return ret
-
-
-class Collection(Packaged):
+class Collection(pathed.Pathed):
     """A collection of templates of a specific type"""
-
-    _default = None
+    _mask = None
 
     def get_path(self, file):
         return os.path.join(self.path, str(file))
 
-    def get_pkg(self, file):
-        return (self.pkg, str(file))
-
 
 class RepositoryCollection(Collection):
-    _default = 'default'
-
+    _mask = '%s'
 
 class ScriptCollection(Collection):
-    _default = 'default.py_tmpl'
+    _mask = '%s.py_tmpl'
 
+class ManageCollection(Collection):
+    _mask = '%s.py_tmpl'
 
-class Template(Packaged):
-    """Finds the paths/packages of various Migrate templates"""
 
-    _repository = 'repository'
-    _script = 'script'
+class Template(pathed.Pathed):
+    """Finds the paths/packages of various Migrate templates.
+    
+    :param path: Templates are loaded from migrate package
+    if `path` is not provided.
+    """
+    pkg = 'migrate.versioning.templates'
     _manage = 'manage.py_tmpl'
 
-    def __init__(self, pkg):
-        super(Template, self).__init__(pkg)
-        self.repository = RepositoryCollection('.'.join((self.pkg,
-                                                         self._repository)))
-        self.script = ScriptCollection('.'.join((self.pkg, self._script)))
-
-    def get_item(self, attr, filename=None, as_pkg=None, as_str=None):
-        item = getattr(self, attr)
-        if filename is None:
-            filename = getattr(item, '_default')
-        if as_pkg:
-            ret = item.get_pkg(filename)
-            if as_str:
-                ret = '.'.join(ret)
-        else:
-            ret = item.get_path(filename)
-        return ret
+    def __new__(cls, path=None):
+        if path is None:
+            path = cls._find_path(cls.pkg)
+        return super(Template, cls).__new__(cls, path)
 
-    def get_repository(self, filename=None, as_pkg=None, as_str=None):
-        return self.get_item('repository', filename, as_pkg, as_str)
-    
-    def get_script(self, filename=None, as_pkg=None, as_str=None):
-        return self.get_item('script', filename, as_pkg, as_str)
+    def __init__(self, path=None):
+        if path is None:
+            path = Template._find_path(self.pkg)
+        super(Template, self).__init__(path)
+        self.repository = RepositoryCollection(os.path.join(path, 'repository'))
+        self.script = ScriptCollection(os.path.join(path, 'script'))
+        self.manage = ManageCollection(os.path.join(path, 'manage'))
 
-    def manage(self, **k):
-        return (self.pkg, self._manage)
+    @classmethod
+    def _find_path(cls, pkg):
+        """Returns absolute path to dotted python package."""
+        tmp_pkg = pkg.rsplit('.', 1)
 
+        if len(tmp_pkg) != 1:
+            return resource_filename(tmp_pkg[0], tmp_pkg[1])
+        else:
+            return resource_filename(tmp_pkg[0], '')
+
+    def _get_item(self, collection, theme=None):
+        """Locates and returns collection.
+        
+        :param collection: name of collection to locate
+        :param type_: type of subfolder in collection (defaults to "_default")
+        :returns: (package, source)
+        :rtype: str, str
+        """
+        item = getattr(self, collection)
+        theme_mask = getattr(item, '_mask')
+        theme = theme_mask % (theme or 'default')
+        return item.get_path(theme)
+
+    def get_repository(self, *a, **kw):
+        """Calls self._get_item('repository', *a, **kw)"""
+        return self._get_item('repository', *a, **kw)
+    
+    def get_script(self, *a, **kw):
+        """Calls self._get_item('script', *a, **kw)"""
+        return self._get_item('script', *a, **kw)
 
-template_pkg = 'migrate.versioning.templates'
-template = Template(template_pkg)
+    def get_manage(self, *a, **kw):
+        """Calls self._get_item('manage', *a, **kw)"""
+        return self._get_item('manage', *a, **kw)
diff --git a/migrate/versioning/version.py b/migrate/versioning/version.py
index d0842e3..ba67e87 100644
--- a/migrate/versioning/version.py
+++ b/migrate/versioning/version.py
@@ -101,7 +101,7 @@ class Collection(pathed.Pathed):
         if os.path.exists(filepath):
             raise Exception('Script already exists: %s' % filepath)
         else:
-            script.PythonScript.create(filepath)
+            script.PythonScript.create(filepath, **k)
 
         self.versions[ver] = Version(ver, self.path, [filename])
         
diff --git a/test/versioning/test_cfgparse.py b/test/versioning/test_cfgparse.py
index 55de7ab..3cd0e56 100644
--- a/test/versioning/test_cfgparse.py
+++ b/test/versioning/test_cfgparse.py
@@ -1,21 +1,27 @@
-from test import fixture
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
 from migrate.versioning import cfgparse
 from migrate.versioning.repository import *
+from migrate.versioning.template import Template
+from test import fixture
+
 
 class TestConfigParser(fixture.Base):
+
     def test_to_dict(self):
         """Correctly interpret config results as dictionaries"""
         parser = cfgparse.Parser(dict(default_value=42))
-        self.assert_(len(parser.sections())==0)
+        self.assert_(len(parser.sections()) == 0)
         parser.add_section('section')
         parser.set('section','option','value')
-        self.assert_(parser.get('section','option')=='value')
-        self.assert_(parser.to_dict()['section']['option']=='value')
+        self.assertEqual(parser.get('section', 'option'), 'value')
+        self.assertEqual(parser.to_dict()['section']['option'], 'value')
     
     def test_table_config(self):
         """We should be able to specify the table to be used with a repository"""
-        default_text=Repository.prepare_config(template.get_repository(as_pkg=True,as_str=True),
-            Repository._config,'repository_name')
-        specified_text=Repository.prepare_config(template.get_repository(as_pkg=True,as_str=True),
-            Repository._config,'repository_name',version_table='_other_table')
-        self.assertNotEquals(default_text,specified_text)
+        default_text = Repository.prepare_config(Template().get_repository(),
+            Repository._config, 'repository_name')
+        specified_text = Repository.prepare_config(Template().get_repository(),
+            Repository._config, 'repository_name', version_table='_other_table')
+        self.assertNotEquals(default_text, specified_text)
diff --git a/test/versioning/test_template.py b/test/versioning/test_template.py
index e92cb95..d0c75a9 100644
--- a/test/versioning/test_template.py
+++ b/test/versioning/test_template.py
@@ -1,17 +1,63 @@
-from test import fixture
-from migrate.versioning.repository import *
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
 import os
+import shutil
 
-class TestPathed(fixture.Base):
+import migrate.versioning.templates
+from migrate.versioning.template import *
+from migrate.versioning import api
+
+from test import fixture
+
+
+class TestTemplate(fixture.Pathed):
     def test_templates(self):
         """We can find the path to all repository templates"""
-        path = str(template)
+        path = str(Template())
         self.assert_(os.path.exists(path))
+
     def test_repository(self):
         """We can find the path to the default repository"""
-        path = template.get_repository()
+        path = Template().get_repository()
         self.assert_(os.path.exists(path))
+
     def test_script(self):
         """We can find the path to the default migration script"""
-        path = template.get_script()
+        path = Template().get_script()
         self.assert_(os.path.exists(path))
+
+    def test_custom_templates_and_themes(self):
+        """Users can define their own templates with themes"""
+        new_templates_dir = os.path.join(self.temp_usable_dir, 'templates')
+        manage_tmpl_file = os.path.join(new_templates_dir, 'manage/custom.py_tmpl')
+        repository_tmpl_file = os.path.join(new_templates_dir, 'repository/custom/README')
+        script_tmpl_file = os.path.join(new_templates_dir, 'script/custom.py_tmpl')
+        MANAGE_CONTENTS = 'print "manage.py"'
+        README_CONTENTS = 'MIGRATE README!'
+        SCRIPT_FILE_CONTENTS = 'print "script.py"'
+        new_repo_dest = self.tmp_repos()
+        new_manage_dest = self.tmp_py()
+
+        # make new templates dir
+        shutil.copytree(migrate.versioning.templates.__path__[0], new_templates_dir)
+        shutil.copytree(os.path.join(new_templates_dir, 'repository/default'),
+            os.path.join(new_templates_dir, 'repository/custom'))
+
+        # edit templates
+        f = open(manage_tmpl_file, 'w').write(MANAGE_CONTENTS)
+        f = open(repository_tmpl_file, 'w').write(README_CONTENTS)
+        f = open(script_tmpl_file, 'w').write(SCRIPT_FILE_CONTENTS)
+
+        # create repository, manage file and python script
+        kw = {}
+        kw['templates_path'] = new_templates_dir
+        kw['templates_theme'] = 'custom'
+        api.create(new_repo_dest, 'repo_name', **kw)
+        api.script('test', new_repo_dest, **kw)
+        api.manage(new_manage_dest, **kw)
+
+        # assert changes
+        self.assertEqual(open(new_manage_dest).read(), MANAGE_CONTENTS)
+        self.assertEqual(open(os.path.join(new_repo_dest, 'README')).read(), README_CONTENTS)
+        self.assertEqual(open(os.path.join(new_repo_dest, 'versions/001_test.py')).read(), SCRIPT_FILE_CONTENTS)