Model generator

grace/cli.py

@@ -5,6 +5,7 @@
 from logging import info, warning, critical
 from click import group, argument, option, pass_context, echo
 from grace.generator import register_generators
+from grace.database import up_migration, down_migration
 from textwrap import dedent
 
 
@@ -46,8 +47,8 @@ def new(ctx, name, database=True):
 def app_cli(ctx, environment):
     app = ctx.obj["app"]
 
-    register_generators(generate)
     app.load(environment)
+    register_generators(generate)
 
 
 @app_cli.group()
@@ -110,6 +111,30 @@ def seed(ctx):
     seed.seed_database()
 
 
+@db.command()
+@argument("revision", default='head')
+@pass_context
+def up(ctx, revision):
+    app = ctx.obj["app"]
+
+    if not app.database_exists:
+        return warning("Database does not exist")
+
+    up_migration(app, revision)
+
+
+@db.command()
+@argument("revision", default='head')
+@pass_context
+def down(ctx, revision):
+    app = ctx.obj["app"]
+
+    if not app.database_exists:
+        return warning("Database does not exist")
+
+    down_migration(app, revision)
+
+
 def _load_database(app):
     if not app.database_exists:
         app.create_database()
@@ -134,4 +159,4 @@ def main():
         from bot import app, bot
         app_cli(obj={"app": app, "bot": bot})
     except ImportError:
-        cli()
+        cli()

grace/config.py

@@ -101,7 +101,7 @@ def database_uri(self) -> Union[str, URL]:
             return self.database.get("url")
 
         return URL.create(
-            self.database.get("adapter"),
+            self.database.get("adapter", "sqlite"),
             self.database.get("user"),
             self.database.get("password"),
             self.database.get("host"),
@@ -139,10 +139,7 @@ def get(self, section_key, value_key, fallback=None) -> Optional[Union[str, int,
     def set_environment(self, environment: str):
         """Set the environment for the configuration.
 
-        :param environment: The environment to set. (Production, Development, Test)
+        :param environment: The environment to set.
         :type environment: str
         """
-        if environment in ["production", "development", "test"]:
-            self.__environment = environment
-        else:
-            raise EnvironmentError("You need to pass a valid environment. [Production, Development, Test]")
+        self.__environment = environment

grace/database.py

@@ -0,0 +1,39 @@
+from alembic.config import Config
+from alembic.command import revision, upgrade, downgrade, show
+from alembic.util.exc import CommandError
+from logging import info, fatal
+
+from alembic.script import ScriptDirectory
+
+
+def generate_migration(app, message):
+    alembic_cfg = Config("alembic.ini")
+    alembic_cfg.config_ini_section = app.config.current_environment
+
+    try:
+        revision(
+            alembic_cfg,
+            message=message,
+            autogenerate=True,
+            sql=False
+        )
+    except CommandError as e:
+        fatal(f"Error creating migration: {e}")
+
+
+def up_migration(app, revision='head'):
+    info(f"Upgrading revision {revision}")
+
+    alembic_cfg = Config("alembic.ini")
+    alembic_cfg.config_ini_section = app.config.current_environment
+
+    upgrade(alembic_cfg, revision=revision)
+
+
+def down_migration(app, revision='head'):
+    info(f"Downgrading revision {revision}")
+
+    alembic_cfg = Config("alembic.ini")
+    alembic_cfg.config_ini_section = app.config.current_environment
+
+    downgrade(alembic_cfg, revision=revision)

grace/generator.py

@@ -22,12 +22,17 @@ def generator() -> Generator:
 ```
 
 """
+import inflect
+
+
 from click import Command, Group
 from pathlib import Path
+from grace.application import Application
 from grace.importer import import_package_modules
 from grace.exceptions import GeneratorError, ValidationError, NoTemplateError
 from cookiecutter.main import cookiecutter
 from jinja2 import Environment, PackageLoader, Template
+from typing import Any
 
 
 def register_generators(command_group: Group):
@@ -69,7 +74,7 @@ class Generator(Command):
     - `NAME`: The name of the generator command (must be defined by subclasses).
     - `OPTIONS`: A dictionary of additional Click options for the command.
     """
-    NAME: str = None
+    NAME: str | None = None
     OPTIONS: dict = {
 
     }
@@ -80,6 +85,8 @@ def __init__(self):
 
         :raises GeneratorError: If the `NAME` attribute is not defined.
         """
+        self.app: Application | None = None
+
         if not self.NAME:
             raise GeneratorError("Generator name must be defined.")
 
@@ -89,6 +96,10 @@ def __init__(self):
     def templates_path(self) -> Path:
         return Path(__file__).parent / 'generators' / 'templates'
 
+    def invoke(self, ctx):
+        self.app = ctx.obj.get("app")
+        return super().invoke(ctx)
+
     def generate(self, *args, **kwargs):
         """Generates template.
 
@@ -108,22 +119,22 @@ def validate(self, *args, **kwargs):
         """Validates the arguments passed to the command."""
         return True
 
-    def generate_template(self, template_dir: str, variables: dict[str, any] = {}):
+    def generate_template(self, template_dir: str, variables: dict[str, Any] = {}):
         """Generates a template using Cookiecutter.
 
         :param template_dir: The name of the template to generate.
         :type template_dir: str
 
         :param variables: The variables to pass to the template. (default is {})
-        :type variables: dict[str, any]
+        :type variables: dict[str, Any]
         """
         template = str(self.templates_path / template_dir)
         cookiecutter(template, extra_context=variables, no_input=True)
 
     def generate_file(
         self,
         template_dir: str,
-        variables: dict[str, any] = {},
+        variables: dict[str, Any] = {},
         output_dir: str = ""
     ):
         """Generate a module using jinja2 template.
@@ -132,17 +143,18 @@ def generate_file(
         :type template_dir: str
 
         :param variables: The variables to pass to the template. (default is {})
-        :type variables: dict[str, any]
+        :type variables: dict[str, Any]
 
         :param output_dir: The output directory for the generated template. (default is None)
         :type output_dir: str
         """
         env = Environment(
-            loader=PackageLoader("grace", self.templates_path / template_dir),
+            loader=PackageLoader('grace', str(self.templates_path / template_dir)),
             extensions=['jinja2_strcase.StrcaseExtension']
         )
 
         env.filters['camel_case_to_space'] = _camel_case_to_space
+        env.filters['pluralize'] = lambda w: inflect.engine().plural(w)
 
         if not env.list_templates():
             raise NoTemplateError(f"No templates found in {template_dir}")
@@ -154,4 +166,4 @@ def generate_file(
         rendered_content = template.render(variables)
 
         with open(f"{output_dir}/{rendered_filename}", "w") as file:
-            file.write(rendered_content)
+            file.write(rendered_content)

grace/generators/migration_generator.py

@@ -0,0 +1,40 @@
+from grace.generator import Generator
+from click.core import Argument
+from logging import info
+from grace.database import generate_migration
+
+
+class MigrationGenerator(Generator):
+    NAME: str = 'migration'
+    OPTIONS: dict = {
+        "params": [
+            Argument(["message"], type=str),
+        ],
+    }
+
+    def generate(self, message: str):
+        """Generates a new Alembic migration using autogenerate.
+
+        Example:
+        ```bash
+        grace generate migration "Add Greeting model"
+        ```
+        """
+        info(f"Generating migration '{message}'")
+        generate_migration(self.app, message)
+
+    def validate(self, message: str, **_kwargs) -> bool:
+        """
+        Validate the migration message.
+
+        The message must be a non-empty string.
+
+        Example:
+        - AddUserModel
+        - add_email_to_users
+        """
+        return bool(message.strip())
+
+
+def generator() -> Generator:
+    return MigrationGenerator()

grace/generators/model_generator.py

@@ -0,0 +1,79 @@
+from grace.generator import Generator
+from re import match
+from logging import info
+from click.core import Argument
+from grace.generators.migration_generator import generate_migration
+
+
+class ModelGenerator(Generator):
+    NAME: str = 'model'
+    OPTIONS: dict = {
+        "params": [
+            Argument(["name"], type=str),
+            Argument(["params"], type=str, nargs=-1)
+        ],
+    }
+
+    def generate(self, name: str, params: tuple[str]):
+        """Generate a new model file along its migration.
+
+        The model will be created in the `bot/models` directory with
+        a SQLAlchemy-style definition. You can specify column names and types
+        during generation using the format `column_name:Type`.
+
+        Supported types are any valid SQLAlchemy column types (e.g., `String`, `Integer`,
+        `Boolean`, etc.). See https://docs.sqlalchemy.org/en/20/core/types.html
+
+        Example:
+        ```bash
+        grace generate model Greeting message:String lang:String
+        ```
+        """
+        info(f"Generating model '{name}'")
+
+        columns, types = self.extract_columns(params)
+        model_columns = map(lambda c: f"{c[0]} = Column({c[1]})", columns)
+
+        self.generate_file(
+            self.NAME,
+            variables={
+                "model_name": name,
+                "model_columns": model_columns,
+                "model_column_types": types
+            },
+            output_dir="bot/models"
+        )
+
+        generate_migration(self.app, f"Create {name}")
+
+    def validate(self, name: str, **_kwargs) -> bool:
+        """Validate the model name.
+
+        A valid model name must:
+        - Start with an uppercase letter.
+        - Contain only letters (A-Z, a-z) and numbers (0-9).
+
+        Examples of valid names:
+        - HelloWorld
+        - User123
+        - ProductItem
+        """
+        return bool(match(r'^[A-Z][a-zA-Z0-9]*$', name))
+
+    def extract_columns(self, params: tuple[str]) -> tuple[list, list]:
+        columns = []
+        types = ['Integer']
+
+        for param in params:
+            name, type = param.split(':')
+
+            if type not in types:
+                types.append(type)
+
+            columns.append((name, type))
+
+        return columns, types
+
+
+def generator() -> Generator:
+    return ModelGenerator()

grace/generators/templates/model/{{ model_name | to_snake }}.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Column, Integer{{ ", {}".format(','.join(model_column_types)) }}
+from bot import app
+from grace.model import Model
+
+
+class {{ model_name | to_camel }}(app.base, Model):
+    __tablename__ = "{{ model_name | to_snake | pluralize }}"
+
+    id = Column(Integer, primary_key=True)
+    {{ model_columns | join('\n    ') }}

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
 	"alembic==1.8.1",
 	"cookiecutter",
 	"jinja2-strcase",
+	"inflect",
 	"mypy",
 	"pytest",
 	"flake8",
@@ -43,3 +44,7 @@ packages = ["grace"]
 
 [tool.mypy]
 exclude = ['grace/generators/templates']
+
+[[tool.mypy.overrides]]
+module = ["jinja2_pluralize.*", "cookiecutter.*"]
+follow_untyped_imports = true