lecrapaud 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

lecrapaud/db/models/base.py

@@ -1,6 +1,181 @@
+"""Base SQLAlchemy model with CRUD operations."""
+
+from functools import wraps
+
 from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy import desc, asc, and_, delete
+from sqlalchemy.inspection import inspect
+from sqlalchemy.orm.attributes import InstrumentedAttribute
+from lecrapaud.db.session import get_db
+
+
+def with_db(func):
+    """Decorator to allow passing an optional db session"""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        db = kwargs.pop("db", None)
+        if db:
+            return func(*args, db=db, **kwargs)
+        with get_db() as db:
+            return func(*args, db=db, **kwargs)
+
+    return wrapper
 
 
 # declarative base class
 class Base(DeclarativeBase):
-    pass
+    @classmethod
+    @with_db
+    def create(cls, db, **kwargs):
+        instance = cls(**kwargs)
+        db.add(instance)
+        db.commit()
+        db.refresh(instance)
+        return instance
+
+    @classmethod
+    @with_db
+    def get(cls, id: int, db=None):
+        return db.get(cls, id)
+
+    @classmethod
+    @with_db
+    def find_by(cls, db=None, **kwargs):
+        return db.query(cls).filter_by(**kwargs).first()
+
+    @classmethod
+    @with_db
+    def get_all(
+        cls, raw=False, db=None, limit: int = 100, order: str = "desc", **kwargs
+    ):
+        order_by_field = (
+            desc(cls.created_at) if order == "desc" else asc(cls.created_at)
+        )
+
+        query = db.query(cls)
+
+        # Apply filters from kwargs
+        for key, value in kwargs.items():
+            if hasattr(cls, key):
+                query = query.filter(getattr(cls, key) == value)
+
+        results = query.order_by(order_by_field).limit(limit).all()
+
+        if raw:
+            return [
+                {
+                    column.name: getattr(row, column.name)
+                    for column in cls.__table__.columns
+                }
+                for row in results
+            ]
+
+        return results
+
+    @classmethod
+    @with_db
+    def filter(cls, db=None, **kwargs):
+        filters = []
+
+        for key, value in kwargs.items():
+            if "__" in key:
+                field, op = key.split("__", 1)
+            else:
+                field, op = key, "eq"
+
+            if not hasattr(cls, field):
+                raise ValueError(f"{field} is not a valid field on {cls.__name__}")
+
+            column: InstrumentedAttribute = getattr(cls, field)
+
+            if op == "eq":
+                filters.append(column == value)
+            elif op == "in":
+                filters.append(column.in_(value))
+            elif op == "gt":
+                filters.append(column > value)
+            elif op == "lt":
+                filters.append(column < value)
+            elif op == "gte":
+                filters.append(column >= value)
+            elif op == "lte":
+                filters.append(column <= value)
+            else:
+                raise ValueError(f"Unsupported operator: {op}")
+
+        return db.query(cls).filter(and_(*filters)).all()
+
+    @classmethod
+    @with_db
+    def update(cls, id: int, db=None, **kwargs):
+        instance = db.get(cls, id)
+        if not instance:
+            return None
+        for key, value in kwargs.items():
+            setattr(instance, key, value)
+        db.commit()
+        db.refresh(instance)
+        return instance
+
+    @classmethod
+    @with_db
+    def upsert(cls, match_fields: list[str], db=None, **kwargs):
+        """
+        Upsert an instance of the model: update if found, else create.
+
+        :param match_fields: list of field names to use for matching
+        :param kwargs: all fields for creation or update
+        """
+        filters = [
+            getattr(cls, field) == kwargs[field]
+            for field in match_fields
+            if field in kwargs
+        ]
+
+        instance = db.query(cls).filter(*filters).first()
+
+        if instance:
+            for key, value in kwargs.items():
+                setattr(instance, key, value)
+        else:
+            instance = cls(**kwargs)
+            db.add(instance)
+
+        db.commit()
+        db.refresh(instance)
+        return instance
+
+    @classmethod
+    @with_db
+    def delete(cls, id: int, db=None):
+        instance = db.get(cls, id)
+        if instance:
+            db.delete(instance)
+            db.commit()
+            return True
+        return False
+
+    @classmethod
+    @with_db
+    def delete_all(cls, db=None, **kwargs):
+        stmt = delete(cls)
+
+        for key, value in kwargs.items():
+            if hasattr(cls, key):
+                stmt = stmt.where(getattr(cls, key) == value)
+
+        db.execute(stmt)
+        db.commit()
+        return True
+
+    @with_db
+    def save(self, db=None):
+        self = db.merge(self)
+        db.add(self)
+        db.commit()
+        db.refresh(self)
+        return self
+
+    def to_json(self):
+        return {c.key: getattr(self, c.key) for c in inspect(self).mapper.column_attrs}

lecrapaud/db/models/dataset.py

@@ -17,9 +17,8 @@ from sqlalchemy import desc, asc, cast, text, func
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 from itertools import chain
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin, with_db
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
 
 # jointures
 dataset_target_association = Table(
@@ -40,7 +39,7 @@ dataset_target_association = Table(
 )
 
 
-class Dataset(Base, CRUDMixin):
+class Dataset(Base):
     __tablename__ = "datasets"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -57,22 +56,22 @@ class Dataset(Base, CRUDMixin):
     path = Column(String(255))  # we do not have this at creation time
     type = Column(String(50), nullable=False)
     size = Column(Integer, nullable=False)
-    train_size = Column(Integer, nullable=False)
+    train_size = Column(Integer)
     val_size = Column(Integer)
-    test_size = Column(Integer, nullable=False)
-    number_of_groups = Column(Integer, nullable=False)
-    list_of_groups = Column(JSON, nullable=False)
+    test_size = Column(Integer)
     corr_threshold = Column(Float, nullable=False)
     max_features = Column(Integer, nullable=False)
     percentile = Column(Float, nullable=False)
-    start_date = Column(DateTime, nullable=False)
-    end_date = Column(DateTime, nullable=False)
-    train_start_date = Column(DateTime, nullable=False)
-    train_end_date = Column(DateTime, nullable=False)
+    number_of_groups = Column(Integer)
+    list_of_groups = Column(JSON)
+    start_date = Column(DateTime)
+    end_date = Column(DateTime)
+    train_start_date = Column(DateTime)
+    train_end_date = Column(DateTime)
     val_start_date = Column(DateTime)
     val_end_date = Column(DateTime)
-    test_start_date = Column(DateTime, nullable=False)
-    test_end_date = Column(DateTime, nullable=False)
+    test_start_date = Column(DateTime)
+    test_end_date = Column(DateTime)
 
     feature_selections = relationship(
         "FeatureSelection",
@@ -111,14 +110,20 @@ class Dataset(Base, CRUDMixin):
         feature = [f.name for f in feature_selection.features]
         return feature
 
-    def get_all_features(self):
+    def get_all_features(self, date_column: str = None, group_column: str = None):
         target_idx = [target.id for target in self.targets]
-        all_features = chain.from_iterable(
-            [f.name for f in fs.features]
-            for fs in self.feature_selections
-            if fs.target_id in target_idx
+        all_features = []
+        if date_column:
+            all_features.append(date_column)
+        if group_column:
+            all_features.append(group_column)
+        all_features += list(
+            chain.from_iterable(
+                [f.name for f in fs.features]
+                for fs in self.feature_selections
+                if fs.target_id in target_idx
+            )
         )
-        all_features = ["DATE", "STOCK"] + list(all_features)
         all_features = list(dict.fromkeys(all_features))
 
         return all_features

lecrapaud/db/models/feature.py

@@ -16,13 +16,12 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.models.feature_selection import feature_selection_association
-from src.db.crud import CRUDMixin
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
+from lecrapaud.db.models.feature_selection import feature_selection_association
 
 
-class Feature(Base, CRUDMixin):
+class Feature(Base):
     __tablename__ = "features"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -42,5 +41,5 @@ class Feature(Base, CRUDMixin):
         "FeatureSelection",
         secondary=feature_selection_association,
         back_populates="features",
-        lazy="selectin"
+        lazy="selectin",
     )

lecrapaud/db/models/feature_selection.py

@@ -24,9 +24,8 @@ from sqlalchemy.orm import (
 )
 from collections.abc import Iterable
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin, with_db
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base, with_db
 
 # jointures
 feature_selection_association = Table(
@@ -47,7 +46,7 @@ feature_selection_association = Table(
 )
 
 
-class FeatureSelection(Base, CRUDMixin):
+class FeatureSelection(Base):
     __tablename__ = "feature_selections"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)

lecrapaud/db/models/feature_selection_rank.py

@@ -18,12 +18,11 @@ from sqlalchemy import desc, asc, cast, text, func
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 from sqlalchemy.dialects.mysql import insert
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin, with_db
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base, with_db
 
 
-class FeatureSelectionRank(Base, CRUDMixin):
+class FeatureSelectionRank(Base):
     __tablename__ = "feature_selection_ranks"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)

lecrapaud/db/models/model.py

@@ -17,12 +17,11 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
 
 
-class Model(Base, CRUDMixin):
+class Model(Base):
     __tablename__ = "models"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)

lecrapaud/db/models/model_selection.py

@@ -17,12 +17,11 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
 
 
-class ModelSelection(Base, CRUDMixin):
+class ModelSelection(Base):
     __tablename__ = "model_selections"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -38,11 +37,20 @@ class ModelSelection(Base, CRUDMixin):
     best_model_params = Column(JSON)
     best_model_path = Column(String(255))
     best_model_id = Column(BigInteger, ForeignKey("models.id", ondelete="CASCADE"))
-    target_id = Column(BigInteger, ForeignKey("targets.id", ondelete="CASCADE"), nullable=False)
-    dataset_id = Column(BigInteger, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
+    target_id = Column(
+        BigInteger, ForeignKey("targets.id", ondelete="CASCADE"), nullable=False
+    )
+    dataset_id = Column(
+        BigInteger, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False
+    )
 
     best_model = relationship("Model", lazy="selectin")
-    model_trainings = relationship("ModelTraining", back_populates="model_selection", cascade="all, delete-orphan", lazy="selectin")
+    model_trainings = relationship(
+        "ModelTraining",
+        back_populates="model_selection",
+        cascade="all, delete-orphan",
+        lazy="selectin",
+    )
     dataset = relationship(
         "Dataset", back_populates="model_selections", lazy="selectin"
     )
@@ -53,4 +61,3 @@ class ModelSelection(Base, CRUDMixin):
             "target_id", "dataset_id", name="uq_model_selection_composite"
         ),
     )
-

lecrapaud/db/models/model_training.py

@@ -17,12 +17,11 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
 
 
-class ModelTraining(Base, CRUDMixin):
+class ModelTraining(Base):
     __tablename__ = "model_trainings"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -40,12 +39,21 @@ class ModelTraining(Base, CRUDMixin):
     training_time = Column(Integer)
     model_id = Column(BigInteger, ForeignKey("models.id"), nullable=False)
     model_selection_id = Column(
-        BigInteger, ForeignKey("model_selections.id", ondelete="CASCADE"), nullable=False
+        BigInteger,
+        ForeignKey("model_selections.id", ondelete="CASCADE"),
+        nullable=False,
     )
 
     model = relationship("Model", lazy="selectin")
-    model_selection = relationship("ModelSelection", back_populates="model_trainings", lazy="selectin")
-    score = relationship("Score", back_populates="model_trainings", cascade="all, delete-orphan", lazy="selectin")
+    model_selection = relationship(
+        "ModelSelection", back_populates="model_trainings", lazy="selectin"
+    )
+    score = relationship(
+        "Score",
+        back_populates="model_trainings",
+        cascade="all, delete-orphan",
+        lazy="selectin",
+    )
 
     __table_args__ = (
         UniqueConstraint(

lecrapaud/db/models/score.py

@@ -16,12 +16,11 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
 
 
-class Score(Base, CRUDMixin):
+class Score(Base):
     __tablename__ = "scores"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -57,6 +56,10 @@ class Score(Base, CRUDMixin):
     threshold = Column(Float)
     precision_at_threshold = Column(Float)
     recall_at_threshold = Column(Float)
-    model_training_id = Column(BigInteger, ForeignKey("model_trainings.id", ondelete="CASCADE"), nullable=False)
+    model_training_id = Column(
+        BigInteger, ForeignKey("model_trainings.id", ondelete="CASCADE"), nullable=False
+    )
 
-    model_trainings = relationship("ModelTraining", back_populates="score", lazy="selectin")
+    model_trainings = relationship(
+        "ModelTraining", back_populates="score", lazy="selectin"
+    )

lecrapaud/db/models/target.py

@@ -17,13 +17,12 @@ from sqlalchemy import desc, asc, cast, text, func
 
 from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
 
-from src.db.setup import get_db
-from src.db.models.base import Base
-from src.db.crud import CRUDMixin
-from src.db.models.dataset import dataset_target_association
+from lecrapaud.db.session import get_db
+from lecrapaud.db.models.base import Base
+from lecrapaud.db.models.dataset import dataset_target_association
 
 
-class Target(Base, CRUDMixin):
+class Target(Base):
     __tablename__ = "targets"
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
@@ -41,13 +40,22 @@ class Target(Base, CRUDMixin):
     description = Column(String(255))
 
     datasets = relationship(
-        "Dataset", secondary=dataset_target_association, back_populates="targets", lazy="selectin"
+        "Dataset",
+        secondary=dataset_target_association,
+        back_populates="targets",
+        lazy="selectin",
     )
     feature_selections = relationship(
-        "FeatureSelection", back_populates="target", cascade="all, delete-orphan", lazy="selectin"
+        "FeatureSelection",
+        back_populates="target",
+        cascade="all, delete-orphan",
+        lazy="selectin",
     )
     model_selections = relationship(
-        "ModelSelection", back_populates="target", cascade="all, delete-orphan", lazy="selectin"
+        "ModelSelection",
+        back_populates="target",
+        cascade="all, delete-orphan",
+        lazy="selectin",
     )
 
     __table_args__ = (

lecrapaud/db/session.py

@@ -0,0 +1,66 @@
+"""Database session management and initialization module."""
+
+from contextlib import contextmanager
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import sessionmaker
+from urllib.parse import urlparse
+from alembic.config import Config
+from alembic import command
+import os
+
+from lecrapaud.config import DB_USER, DB_PASSWORD, DB_HOST, DB_PORT, DB_NAME, DB_URI
+
+_engine = None
+_SessionLocal = None
+DATABASE_URL = (
+    f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}" or DB_URI
+)
+
+
+def init_db(uri: str = None):
+    global _engine, _SessionLocal
+
+    uri = uri if uri else DATABASE_URL
+    # Extract DB name from URI to connect without it
+    parsed = urlparse(uri)
+    db_name = parsed.path.lstrip("/")  # remove leading slash
+
+    # Build root engine (no database in URI)
+    root_uri = uri.replace(f"/{db_name}", "/")
+
+    # Step 1: Connect to MySQL without a database
+    root_engine = create_engine(root_uri)
+
+    # Step 2: Create database if it doesn't exist
+    with root_engine.connect() as conn:
+        conn.execute(text(f"CREATE DATABASE IF NOT EXISTS {DB_NAME}"))
+        conn.commit()
+
+    # Step 3: Connect to the newly created database
+    _engine = create_engine(uri, echo=False)
+
+    # Step 4: Create session factory
+    _SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=_engine)
+
+    # Step 5: Apply Alembic migrations programmatically
+    project_root = os.path.abspath(os.path.dirname(__file__))
+    alembic_cfg_path = os.path.join(project_root, "alembic.ini")
+
+    alembic_cfg = Config(alembic_cfg_path)
+    alembic_cfg.set_main_option("sqlalchemy.url", uri or os.getenv("DATABASE_URL"))
+    command.upgrade(alembic_cfg, "head")
+
+
+# Dependency to get a session instance
+@contextmanager
+def get_db():
+    if _SessionLocal is None:
+        raise RuntimeError("Database not initialized. Call `init_db()` first.")
+    db = _SessionLocal()
+    try:
+        yield db
+    except Exception as e:
+        db.rollback()
+        raise Exception(e) from e
+    finally:
+        db.close()

lecrapaud/experiment.py

@@ -0,0 +1,64 @@
+import pandas as pd
+import os
+from pathlib import Path
+
+os.environ["COVERAGE_FILE"] = str(Path(".coverage").resolve())
+
+# Internal
+from lecrapaud.directories import tmp_dir
+from lecrapaud.utils import logger
+from lecrapaud.config import PYTHON_ENV
+from lecrapaud.db import (
+    Dataset,
+    Target,
+)
+from lecrapaud.db.session import get_db
+
+
+def create_dataset(
+    data: pd.DataFrame,
+    corr_threshold,
+    percentile,
+    max_features,
+    date_column,
+    group_column,
+    **kwargs,
+):
+    dates = {}
+    if date_column:
+        dates["start_date"] = pd.to_datetime(data[date_column].iat[0])
+        dates["end_date"] = pd.to_datetime(data[date_column].iat[-1])
+
+    groups = {}
+    if group_column:
+        groups["number_of_groups"] = data[group_column].nunique()
+        groups["list_of_groups"] = data[group_column].unique().tolist()
+
+    with get_db() as db:
+        all_targets = Target.get_all(db=db)
+        targets = [target for target in all_targets if target.name in data.columns]
+        dataset_name = f"data_{groups["number_of_groups"] if group_column else 'ng'}_{corr_threshold}_{percentile}_{max_features}_{dates['start_date'].date() if date_column else 'nd'}_{dates['end_date'].date() if date_column else 'nd'}"
+
+        dataset_dir = f"{tmp_dir}/{dataset_name}"
+        preprocessing_dir = f"{dataset_dir}/preprocessing"
+        data_dir = f"{dataset_dir}/data"
+        os.makedirs(dataset_dir, exist_ok=True)
+        os.makedirs(preprocessing_dir, exist_ok=True)
+        os.makedirs(data_dir, exist_ok=True)
+
+        dataset = Dataset.upsert(
+            match_fields=["name"],
+            db=db,
+            name=dataset_name,
+            path=Path(dataset_dir).resolve(),
+            type="training",
+            size=data.shape[0],
+            corr_threshold=corr_threshold,
+            percentile=percentile,
+            max_features=max_features,
+            **groups,
+            **dates,
+            targets=targets,
+        )
+
+        return dataset