Merge pull request #166 from spacemanspiff2007/dev

Added config option to specify source folder

README.md

@@ -26,10 +26,12 @@ Just install from pypi:
 Usage: aerich [OPTIONS] COMMAND [ARGS]...
 
 Options:
+  -V, --version      Show the version and exit.
   -c, --config TEXT  Config file.  [default: aerich.ini]
-  --app TEXT         Tortoise-ORM app name.  [default: models]
+  --app TEXT         Tortoise-ORM app name.
   -n, --name TEXT    Name of section in .ini file to use for aerich config.
                      [default: aerich]
+
   -h, --help         Show this message and exit.
 
 Commands:
@@ -40,7 +42,7 @@ Commands:
   init-db    Generate schema and generate app migrate location.
   inspectdb  Introspects the database tables to standard output as...
   migrate    Generate migrate changes file.
-  upgrade    Upgrade to latest version.
+  upgrade    Upgrade to specified version.
 ```
 
 ## Usage
@@ -68,10 +70,12 @@ Usage: aerich init [OPTIONS]
 
   Init config file and generate root migrate location.
 
-Options:
+OOptions:
-  -t, --tortoise-orm TEXT  Tortoise-ORM config module dict variable, like settings.TORTOISE_ORM.
+  -t, --tortoise-orm TEXT  Tortoise-ORM config module dict variable, like
-                           [required]
+                           settings.TORTOISE_ORM.  [required]
+
   --location TEXT          Migrate store location.  [default: ./migrations]
+  -s, --src_folder TEXT    Folder of the source, relative to the project root.
   -h, --help               Show this message and exit.
 ```
 

aerich/cli.py

@@ -1,6 +1,5 @@
 import asyncio
 import os
-import sys
 from configparser import ConfigParser
 from functools import wraps
 from pathlib import Path
@@ -16,6 +15,7 @@ from tortoise.utils import get_schema_sql
 from aerich.inspectdb import InspectDb
 from aerich.migrate import Migrate
 from aerich.utils import (
+    add_src_path,
     get_app_connection,
     get_app_connection_name,
     get_models_describe,
@@ -23,7 +23,6 @@ from aerich.utils import (
     get_version_content_from_file,
     write_version_file,
 )
-
 from . import __version__
 from .enums import Color
 from .models import Aerich
@@ -74,6 +73,10 @@ async def cli(ctx: Context, config, app, name):
 
         location = parser[name]["location"]
         tortoise_orm = parser[name]["tortoise_orm"]
+        src_folder = parser[name]["src_folder"]
+
+        # Add specified source folder to path
+        add_src_path(src_folder)
 
         tortoise_config = get_tortoise_config(ctx, tortoise_orm)
         app = app or list(tortoise_config.get("apps").keys())[0]
@@ -214,19 +217,35 @@ async def history(ctx: Context):
 @click.option(
     "--location", default="./migrations", show_default=True, help="Migrate store location.",
 )
+@click.option(
+    "-s",
+    "--src_folder",
+    default=".",
+    show_default=False,
+    help="Folder of the source, relative to the project root.",
+)
 @click.pass_context
 @coro
-async def init(
+async def init(ctx: Context, tortoise_orm, location, src_folder):
-    ctx: Context, tortoise_orm, location,
-):
     config_file = ctx.obj["config_file"]
     name = ctx.obj["name"]
     if Path(config_file).exists():
-        return click.secho("You have inited", fg=Color.yellow)
+        return click.secho("Configuration file already created", fg=Color.yellow)
+
+    if os.path.isabs(src_folder):
+        src_folder = os.path.relpath(os.getcwd(), src_folder)
+    # Add ./ so it's clear that this is relative path
+    if not src_folder.startswith("./"):
+        src_folder = "./" + src_folder
+
+    # check that we can find the configuration, if not we can fail before the config file gets created
+    add_src_path(src_folder)
+    get_tortoise_config(ctx, tortoise_orm)
 
     parser.add_section(name)
     parser.set(name, "tortoise_orm", tortoise_orm)
     parser.set(name, "location", location)
+    parser.set(name, "src_folder", src_folder)
 
     with open(config_file, "w", encoding="utf-8") as f:
         parser.write(f)
@@ -294,7 +313,6 @@ async def inspectdb(ctx: Context, table: List[str]):
 
 
 def main():
-    sys.path.insert(0, ".")
     cli()
 
 

aerich/utils.py

@@ -1,12 +1,30 @@
 import importlib
+import os
 import re
+import sys
 from pathlib import Path
 from typing import Dict
 
-from click import BadOptionUsage, Context
+from click import BadOptionUsage, Context, ClickException
 from tortoise import BaseDBAsyncClient, Tortoise
 
 
+def add_src_path(path: str) -> str:
+    """
+    add a folder to the paths so we can import from there
+    :param path: path to add
+    :return: absolute path
+    """
+    if not os.path.isabs(path):
+        # use the absolute path, otherwise some other things (e.g. __file__) won't work properly
+        path = os.path.abspath(path)
+    if not os.path.isdir(path):
+        raise ClickException(f"Specified source folder does not exist: {path}")
+    if path not in sys.path:
+        sys.path.insert(0, path)
+    return path
+
+
 def get_app_connection_name(config, app_name: str) -> str:
     """
     get connection name
@@ -42,7 +60,12 @@ def get_tortoise_config(ctx: Context, tortoise_orm: str) -> dict:
     splits = tortoise_orm.split(".")
     config_path = ".".join(splits[:-1])
     tortoise_config = splits[-1]
-    config_module = importlib.import_module(config_path)
+
+    try:
+        config_module = importlib.import_module(config_path)
+    except ModuleNotFoundError as e:
+        raise ClickException(f"Error while importing configuration module: {e}") from None
+
     config = getattr(config_module, tortoise_config, None)
     if not config:
         raise BadOptionUsage(