nomy-data-models 0.2.4

#!/usr/bin/env python
"""Script to test the conversion of Python models to Rust models.

This script tests the functionality of the py_to_rust module by:
1. Finding all SQLAlchemy models in the nomy_data_models package
2. Generating Rust code for each model
3. Finding all enums in the nomy_data_models package
4. Generating Rust code for each enum
5. Writing the generated code to files in the test_output directory
"""

import inspect
import os
import sys
from enum import Enum
from pathlib import Path
from typing import Dict, Type, Tuple

# 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 (
    get_all_models,
    generate_rust_model,
    generate_rust_enum,
    generate_rust_models,
)
from nomy_data_models import models


def test_get_all_models() -> Dict[str, Type]:
    """Test the get_all_models function.

    Returns:
        Dict[str, Type]: Dictionary of model name to model class
    """
    print("Testing get_all_models...")
    all_models = get_all_models()
    print(f"Found {len(all_models)} models:")
    for name in all_models:
        print(f"  - {name}")
    return all_models


def test_generate_rust_model(model_name: str, model_class: Type) -> str:
    """Test the generate_rust_model function for a specific model.

    Args:
        model_name: Name of the model
        model_class: Model class

    Returns:
        str: Generated Rust code
    """
    print(f"\nTesting generate_rust_model for {model_name}...")
    rust_code = generate_rust_model(model_class)
    print(f"Generated {len(rust_code.splitlines())} lines of Rust code")
    print("First 10 lines:")
    for line in rust_code.splitlines()[:10]:
        print(f"  {line}")
    return rust_code


def test_generate_rust_enum() -> Dict[str, Type[Enum]]:
    """Test the generate_rust_enum function.

    Returns:
        Dict[str, Type[Enum]]: Dictionary of enum name to enum class
    """
    print("\nTesting generate_rust_enum...")
    # Find all enums in the models module
    enums = {}

    for name in dir(models):
        item = getattr(models, name)
        if (
            inspect.isclass(item)
            and issubclass(item, Enum)
            and item != Enum
            and item.__module__.startswith("nomy_data_models.models")
        ):
            enums[name] = item

    print(f"Found {len(enums)} enums:")
    for name in enums:
        print(f"  - {name}")

    if enums:
        # Test the first enum
        name, enum_class = next(iter(enums.items()))
        rust_code = generate_rust_enum(enum_class)
        print(f"\nGenerated Rust code for {name}:")
        for line in rust_code.splitlines()[:10]:
            print(f"  {line}")

    return enums


def test_generate_rust_models() -> Path:
    """Test the generate_rust_models function.

    Returns:
        Path: Path to the output directory
    """
    print("\nTesting generate_rust_models...")
    # Create a temporary directory for the output
    output_dir = Path(__file__).parent / "test_output" / "models"
    os.makedirs(output_dir, exist_ok=True)

    # Generate the Rust models
    generate_rust_models(str(output_dir))

    # Check the output
    print(f"Generated Rust models in {output_dir}")
    files = list(output_dir.glob("*.rs"))
    print(f"Found {len(files)} Rust files:")
    for file in files:
        print(f"  - {file.name}")

    # Check the mod.rs file
    mod_file = output_dir / "mod.rs"
    if mod_file.exists():
        print("\nContents of mod.rs:")
        with open(mod_file, "r") as f:
            for line in f.readlines()[:10]:
                print(f"  {line.strip()}")

    return output_dir


def main() -> None:
    """Main entry point for the script."""
    print("Testing Python to Rust conversion...\n")

    # Test get_all_models
    all_models = test_get_all_models()

    # Test generate_rust_model for each model
    for model_name, model_class in all_models.items():
        test_generate_rust_model(model_name, model_class)

    # Test generate_rust_enum
    test_generate_rust_enum()

    # Test generate_rust_models
    test_generate_rust_models()

    print("\nAll tests completed successfully!")


if __name__ == "__main__":
    main()