- more fixes for #163, make sure env.py load figures out to load env.pyc or

env.pyo as well

3 files changed, 45 insertions, 12 deletions

diff --git a/alembic/util.py b/alembic/util.py
index 93b6b76..4bcd8b5 100644
--- a/alembic/util.py
+++ b/alembic/util.py
@@ -198,13 +198,28 @@ def load_python_file(dir_, filename):
     path = os.path.join(dir_, filename)
     _, ext = os.path.splitext(filename)
     if ext == ".py":
-        module = load_module_py(module_id, path)
+        if os.path.exists(path):
+            module = load_module_py(module_id, path)
+        elif os.path.exists(simple_pyc_file_from_path(path)):
+            # look for sourceless load
+            module = load_module_pyc(module_id, simple_pyc_file_from_path(path))
+        else:
+            raise ImportError("Can't find Python file %s" % path)
     elif ext in (".pyc", ".pyo"):
         module = load_module_pyc(module_id, path)
     del sys.modules[module_id]
     return module
 
 def simple_pyc_file_from_path(path):
+    """Given a python source path, return the so-called
+    "sourceless" .pyc or .pyo path.
+
+    This just a .pyc or .pyo file where the .py file would be.
+
+    Even with PEP-3147, which normally puts .pyc/.pyo files in __pycache__,
+    this use case remains supported as a so-called "sourceless module import".
+
+    """
     if sys.flags.optimize:
         return path + "o"  # e.g. .pyo
     else:
diff --git a/tests/__init__.py b/tests/__init__.py
index ad5b033..9bb40f1 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -303,7 +303,7 @@ def _testing_config():
     return Config(os.path.join(staging_directory, 'test_alembic.ini'))
 
 
-def staging_env(create=True, template="generic"):
+def staging_env(create=True, template="generic", sourceless=False):
     from alembic import command, script
     cfg = _testing_config()
     if create:
@@ -311,6 +311,19 @@ def staging_env(create=True, template="generic"):
         if os.path.exists(path):
             shutil.rmtree(path)
         command.init(cfg, path)
+        if sourceless:
+            try:
+                # do an import so that a .pyc/.pyo is generated.
+                util.load_python_file(path, 'env.py')
+            except AttributeError:
+                # we don't have the migration context set up yet
+                # so running the .env py throws this exception.
+                # theoretically we could be using py_compiler here to
+                # generate .pyc/.pyo without importing but not really
+                # worth it.
+                pass
+            make_sourceless(os.path.join(path, "env.py"))
+
     sc = script.ScriptDirectory.from_config(cfg)
     return sc
 
@@ -339,16 +352,21 @@ def write_script(scriptdir, rev_id, content, encoding='ascii', sourceless=False)
     script.nextrev = old.nextrev
 
     if sourceless:
-        # note that if -O is set, you'd see pyo files here,
-        # the pyc util function looks at sys.flags.optimize to handle this
-        assert os.access(pyc_path, os.F_OK)
-        # look for a non-pep3147 path here.
-        # if not present, need to copy from __pycache__
-        simple_pyc_path = util.simple_pyc_file_from_path(path)
-        if not os.access(simple_pyc_path, os.F_OK):
-            shutil.copyfile(pyc_path, simple_pyc_path)
-        os.unlink(path)
+        make_sourceless(path)
+
+def make_sourceless(path):
+    # note that if -O is set, you'd see pyo files here,
+    # the pyc util function looks at sys.flags.optimize to handle this
+    pyc_path = util.pyc_file_from_path(path)
+    assert os.access(pyc_path, os.F_OK)
+
+    # look for a non-pep3147 path here.
+    # if not present, need to copy from __pycache__
+    simple_pyc_path = util.simple_pyc_file_from_path(path)
 
+    if not os.access(simple_pyc_path, os.F_OK):
+        shutil.copyfile(pyc_path, simple_pyc_path)
+    os.unlink(path)
 
 def three_rev_fixture(cfg):
     a = util.rev_id()
diff --git a/tests/test_versioning.py b/tests/test_versioning.py
index a4be95f..41e4be1 100644
--- a/tests/test_versioning.py
+++ b/tests/test_versioning.py
@@ -99,7 +99,7 @@ class VersioningTest(unittest.TestCase):
 
     @classmethod
     def setup_class(cls):
-        cls.env = staging_env()
+        cls.env = staging_env(sourceless=cls.sourceless)
         cls.cfg = _sqlite_testing_config()
 
     @classmethod