nomy-data-models 0.2.4

#!/usr/bin/env python
"""Script to generate Rust models from Python SQLAlchemy models.

This script is a command-line interface to the py_to_rust module. It generates
Rust models from Python SQLAlchemy models and writes them to the specified
output directory.

Usage:
    python generate_rust.py [output_dir] [--verify]

If output_dir is not specified, the default is src/models.
The --verify flag will check if the generated Rust models build correctly.
"""

import sys
import subprocess
import argparse
import inspect
import logging
from pathlib import Path
from typing import List, Type, Dict
from enum import Enum

# Add the parent directory to the path so we can import the package
sys.path.insert(0, str(Path(__file__).parent.parent))

from nomy_data_models.py_to_rust import (
    generate_rust_enum,
    generate_rust_model,
)
import nomy_data_models.models as models_module
from nomy_data_models.models.base import BaseModel
from nomy_data_models.utils.string import to_snake_case


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


def get_all_enums() -> List[Type[Enum]]:
    """Get all enum classes from the models.enums module.

    Returns:
        List[Type[Enum]]: List of enum classes
    """
    from nomy_data_models.models import enums as enums_module

    enum_classes = []
    for _, obj in inspect.getmembers(enums_module):
        if (
            inspect.isclass(obj)
            and issubclass(obj, Enum)
            and obj.__module__ == enums_module.__name__
        ):
            enum_classes.append(obj)

    return enum_classes


def get_concrete_models() -> Dict[str, Type[BaseModel]]:
    """Get all concrete model classes from the models module.

    This function detects models that should be concrete by examining
    if they explicitly set __abstract__ = False in their class definition.

    Returns:
        Dict[str, Type[BaseModel]]: Dictionary of model name to model class
    """
    result = {}

    # Get all classes from the models module
    for name, obj in inspect.getmembers(models_module):
        # Check if it's a class that inherits from BaseModel
        if (
            inspect.isclass(obj)
            and issubclass(obj, BaseModel)
            and obj != BaseModel
            and obj.__module__.startswith("nomy_data_models.models")
        ):

            # Check if __abstract__ is explicitly defined in the class itself (not inherited)
            # We use obj.__dict__ to check only attributes defined directly in the class
            if "__abstract__" in obj.__dict__ and obj.__dict__["__abstract__"] is False:
                logger.info(f"Including concrete class {name} (__abstract__ = False)")
                result[name] = obj
            else:
                logger.info(
                    f"Skipping class {name} (not explicitly marked as concrete)"
                )

    return result


def write_file(path: Path, content: str) -> None:
    """Write content to a file, creating parent directories if needed.

    Args:
        path: Path to the file
        content: Content to write
    """
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            f.write(content)
        logger.debug(f"Wrote {len(content)} bytes to {path}")
    except Exception as e:
        logger.error(f"Failed to write to {path}: {e}")
        raise


def verify_rust_models(output_dir: str) -> bool:
    """Verify that the generated Rust models build correctly.

    Args:
        output_dir: Directory containing the generated Rust models

    Returns:
        bool: True if verification succeeded, False otherwise
    """
    logger.info("Verifying Rust models build correctly...")

    try:
        # Create a temporary directory for verification
        verify_dir = Path("target/rust_verify")
        verify_dir.mkdir(parents=True, exist_ok=True)

        # Create a simple Cargo.toml file
        cargo_toml_path = verify_dir / "Cargo.toml"
        with open(cargo_toml_path, "w") as f:
            f.write(
                """
[package]
name = "nomy-models"
version = "0.1.0"
edition = "2021"

[dependencies]
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
chrono = { version = "0.4", features = ["serde"] }
uuid = { version = "1.3", features = ["serde", "v4"] }
rust_decimal = { version = "1.29", features = ["serde"] }
"""
            )

        # Create a simple lib.rs file
        src_dir = verify_dir / "src"
        src_dir.mkdir(exist_ok=True)

        with open(src_dir / "lib.rs", "w") as f:
            f.write(
                """
pub mod models;
pub mod enums;

pub use models::*;
pub use enums::*;
"""
            )

        # Create models directory
        models_dir = src_dir / "models"
        models_dir.mkdir(exist_ok=True)

        # Create enums directory
        enums_dir = src_dir / "enums"
        enums_dir.mkdir(exist_ok=True)

        # Copy all generated model files to the models directory
        for file_path in Path("src/models").glob("*.rs"):
            with open(file_path, "r") as src_file:
                content = src_file.read()

            with open(models_dir / file_path.name, "w") as dest_file:
                dest_file.write(content)

        # Copy all generated enum files to the enums directory
        for file_path in Path("src/enums").glob("*.rs"):
            with open(file_path, "r") as src_file:
                content = src_file.read()

            with open(enums_dir / file_path.name, "w") as dest_file:
                dest_file.write(content)

        # Try to build the Rust project
        result = subprocess.run(
            ["cargo", "check", "--manifest-path", str(cargo_toml_path)],
            capture_output=True,
            text=True,
        )

        if result.returncode != 0:
            logger.error(f"Failed to build Rust models:\n{result.stderr}")
            return False

        logger.info("Successfully verified Rust models build correctly.")
        return True
    except Exception as e:
        logger.error(f"Failed to verify Rust models: {e}")
        return False


def main() -> int:
    """Generate Rust models from Python SQLAlchemy models.

    Returns:
        int: Exit code (0 for success, 1 for failure)
    """
    parser = argparse.ArgumentParser(
        description="Generate Rust models from Python SQLAlchemy models"
    )
    parser.add_argument(
        "output_dir",
        nargs="?",
        default="src/models",
        help="Directory to write Rust models to",
    )
    parser.add_argument(
        "--verify",
        action="store_true",
        help="Verify that the generated Rust models build correctly",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )
    args = parser.parse_args()

    # Set log level based on verbosity
    if args.verbose:
        logger.setLevel(logging.DEBUG)

    output_dir = args.output_dir
    verify = args.verify

    logger.info(f"Generating Rust models in {output_dir}...")

    try:
        # Create output directory if it doesn't exist
        Path(output_dir).mkdir(parents=True, exist_ok=True)

        models_dir = Path("src/models")
        enums_dir = Path("src/enums")
        models_dir.mkdir(exist_ok=True)
        enums_dir.mkdir(exist_ok=True)

        # Get all concrete models
        all_models = get_concrete_models()

        # Get all enum classes
        all_enums = get_all_enums()

        logger.info(f"Found {len(all_models)} models and {len(all_enums)} enums")

        # Initialize mod.rs content lists
        models_mod_content = [
            "//! Model definitions for Nomy wallet analysis data processing.",
            "//!",
            "//! This file is generated automatically from the Python models.",
            "//! Do not edit this file manually.",
            "",
        ]

        enums_mod_content = [
            "//! Enum definitions for Nomy wallet analysis data processing.",
            "//!",
            "//! This file is generated automatically from the Python enums.",
            "//! Do not edit this file manually.",
            "",
        ]

        # Track generated models
        generated_models: List[str] = []

        # Generate Rust models
        for model_name, model_class in all_models.items():
            logger.info(f"Generating Rust model for {model_name}")

            # Generate Rust model
            rust_code = generate_rust_model(model_class)
            if not rust_code:
                logger.warning(f"Failed to generate Rust model for {model_name}")
                continue

            # Write to file
            file_name = to_snake_case(model_name) + ".rs"
            write_file(models_dir / file_name, rust_code)

            # Add to models/mod.rs
            snake_case_name = to_snake_case(model_name)
            models_mod_content.append(f"pub mod {snake_case_name};")
            models_mod_content.append(f"pub use {snake_case_name}::{model_name};")
            models_mod_content.append("")

            # Track generated model
            generated_models.append(model_name)

        # Write models/mod.rs
        write_file(models_dir / "mod.rs", "\n".join(models_mod_content))

        # Generate Rust enums
        for enum_class in all_enums:
            enum_name = enum_class.__name__
            logger.info(f"Generating Rust enum for {enum_name}")

            # Generate Rust enum
            rust_code = generate_rust_enum(enum_class)
            if not rust_code:
                logger.warning(f"Failed to generate Rust enum for {enum_name}")
                continue

            # Write to file
            file_name = to_snake_case(enum_name) + ".rs"
            write_file(enums_dir / file_name, rust_code)

            # Add to enums/mod.rs
            snake_case_name = to_snake_case(enum_name)
            enums_mod_content.append(f"pub mod {snake_case_name};")
            enums_mod_content.append(f"pub use {snake_case_name}::{enum_name};")
            enums_mod_content.append("")

        # Write enums/mod.rs
        write_file(enums_dir / "mod.rs", "\n".join(enums_mod_content))

        # Generate lib.rs
        lib_rs_content = [
            "//! Nomy Data Models",
            "//!",
            "//! This crate provides data model definitions for Nomy wallet analysis data processing.",
            "//! These models are shared across multiple services and are generated from Python SQLAlchemy models.",
            "",
            "pub mod models;",
            "pub mod enums;",
            "",
            "/// Re-export all models for convenience",
            "pub use models::*;",
            "pub use enums::*;",
            "",
            "/// Error types for the crate",
            "pub mod error {",
            "    use thiserror::Error;",
            "",
            "    /// Error type for Nomy Data Models",
            "    #[derive(Error, Debug)]",
            "    pub enum NomyDataModelError {",
            "        /// Error when serializing or deserializing data",
            '        #[error("Serialization error: {0}")]',
            "        SerializationError(#[from] serde_json::Error),",
            "",
            "        /// Error when parsing a date or time",
            '        #[error("Date/time parsing error: {0}")]',
            "        DateTimeError(#[from] chrono::ParseError),",
            "",
            "        /// Other errors",
            '        #[error("Other error: {0}")]',
            "        Other(String),",
            "    }",
            "}",
            "",
            "/// Result type for the crate",
            "pub type Result<T> = std::result::Result<T, error::NomyDataModelError>;",
            "",
            "/// Version of the crate",
            'pub const VERSION: &str = env!("CARGO_PKG_VERSION");',
        ]

        # Create src directory if it doesn't exist
        src_dir = Path("src")
        src_dir.mkdir(exist_ok=True)

        # Write lib.rs
        write_file(src_dir / "lib.rs", "\n".join(lib_rs_content))

        # Update imports in model files to use crate::enums instead of crate::models
        for file_path in models_dir.glob("*.rs"):
            with open(file_path, "r") as f:
                content = f.read()

            # Update imports to use crate::enums instead of crate::models
            content = content.replace("use crate::models::", "use crate::enums::")

            with open(file_path, "w") as f:
                f.write(content)

        # Check if all models were generated
        missing_models = [
            name for name in all_models.keys() if name not in generated_models
        ]
        if missing_models:
            logger.error(
                f"The following models were not generated: {', '.join(missing_models)}"
            )
            return 1

        logger.info(
            f"Successfully generated {len(generated_models)} Rust models and {len(all_enums)} enums."
        )

        # Verify that the generated Rust models build correctly
        if verify and not verify_rust_models(output_dir):
            return 1

        logger.info("Done!")
        return 0
    except Exception as e:
        logger.error(f"An error occurred: {e}", exc_info=True)
        return 1


if __name__ == "__main__":
    sys.exit(main())