make cli more user friendly by using enums and types for input instead of string

atlas_provider_sqlalchemy/main.py

@@ -1,5 +1,7 @@
 import os
 from pathlib import Path
+from typing import Optional
+from enum import Enum
 import typer
 from sqlalchemy.orm import DeclarativeBase
 from typing import Type
@@ -8,14 +10,25 @@
 app = typer.Typer(no_args_is_help=True)
 
 
-def run(dialect: str, modles_dir: str = '', debug: bool = False) -> Type[DeclarativeBase]:
-    models_dir = Path(modles_dir) or Path(os.getcwd())
-    base = get_declarative_base(models_dir, debug)
-    return dump_ddl(dialect, base)
+class Dialect(str, Enum):
+    mysql = "mysql"
+    mariadb = "mariadb"
+    postgresql = "postgresql"
+    sqlite = "sqlite"
+    mssql = "mssql"
+
+
+def run(dialect: Dialect, path: Path, debug: bool = False) -> Type[DeclarativeBase]:
+    base = get_declarative_base(path, debug)
+    return dump_ddl(dialect.value, base)
 
 
 @app.command()
-def load(dialect: str = typer.Option(default=...), path: str = '', debug: bool = False):
+def load(dialect: Dialect = Dialect.mysql,
+         path: Optional[Path] = typer.Option(None, exists=True, help="Path to directory of the sqlalchemy models."),
+         debug: bool = True):
+    if path is None:
+        path = Path(os.getcwd())
     run(dialect, path, debug)
 
 

tests/test_main.py

@@ -1,17 +1,14 @@
 from atlas_provider_sqlalchemy.ddl import ModelsNotFoundError
-from atlas_provider_sqlalchemy.main import run, get_declarative_base
+from atlas_provider_sqlalchemy.main import run, get_declarative_base, Dialect
 from pathlib import Path
 from sqlalchemy.orm import DeclarativeBase
 import pytest
 
-POSTGRES_DIALECT = 'postgresql+psycopg2'
-MYSQL_DIALECT = 'mysql+pymysql'
-
 
 def test_run_postgres(capsys):
     with open('tests/ddl_postgres.sql', 'r') as f:
         expected_ddl = f.read()
-    base = run(POSTGRES_DIALECT, "tests/models")
+    base = run(Dialect.postgresql, Path("tests/models"))
     captured = capsys.readouterr()
     assert captured.out == expected_ddl
     base.metadata.clear()
@@ -20,7 +17,7 @@ def test_run_postgres(capsys):
 def test_run_mysql(capsys):
     with open('tests/ddl_mysql.sql', 'r') as f:
         expected_ddl = f.read()
-    base = run(MYSQL_DIALECT, "tests/models")
+    base = run(Dialect.mysql, Path("tests/models"))
     captured = capsys.readouterr()
     assert captured.out == expected_ddl
     base.metadata.clear()