codeflash/code_to_optimize/book_catalog.py

from __future__ import annotations

from time import time
from typing import Any, cast

from _typeshed import SupportsDunderGT, SupportsDunderLT
from sqlalchemy import Boolean, Column, ForeignKey, Integer, Text
from sqlalchemy.engine import Engine, create_engine
from sqlalchemy.orm import DeclarativeBase, Session, relationship, sessionmaker
from sqlalchemy.orm.relationships import _RelationshipDeclared

POSTGRES_CONNECTION_STRING: str = "postgresql://cf_developer:XJcbU37MBYeh4dDK6PTV5n@sqlalchemy-experiments.postgres.database.azure.com:5432/postgres"


class Base(DeclarativeBase):
    pass


class Author(Base):
    __tablename__: str = "authors"

    id: Column[int] = Column(Integer, primary_key=True)
    name: Column[str] = Column(Text, nullable=False)


class Book(Base):
    __tablename__: str = "books"

    id: Column[int] = Column(Integer, primary_key=True)
    title: Column[str] = Column(Text, nullable=False)
    author_id: Column[int] = Column(Integer, ForeignKey("authors.id"), nullable=False)
    is_bestseller: Column[bool] = Column(Boolean, default=False)

    author: _RelationshipDeclared[Author] = relationship("Author", backref="books")


def init_table() -> Session:
    engine: Engine = create_engine(POSTGRES_CONNECTION_STRING, echo=True)
    session: Session = sessionmaker(bind=engine)()
    i: int
    for i in range(50):
        author: Author = Author(id=i, name=f"author{i}")
        session.add(author)
    for i in range(100000):
        book: Book = Book(id=i, title=f"book{i}", author_id=i % 50, is_bestseller=i % 2 == 0)
        session.add(book)
    session.commit()

    return session


def get_authors(session: Session) -> list[Author]:
    books: list[Book] = session.query(Book).all()
    _authors: list[Author] = []
    book: Book
    for book in books:
        _authors.append(book.author)
    return sorted(
        list(set(_authors)),
        key=lambda x: cast(SupportsDunderLT[Any] | SupportsDunderGT[Any], x.id),
    )


if __name__ == "__main__":
    # _session = init_table()
    engine: Engine = create_engine(POSTGRES_CONNECTION_STRING, echo=True)
    session_factory: sessionmaker[Session] = sessionmaker(bind=engine)
    _session: Session = session_factory()
    _t: float = time()
    authors: list[Author] = get_authors(_session)
    print("TIME TAKEN", time() - _t)
    authors_name = list(map(lambda x: x.name, authors))
    print("len(authors_name)", len(authors_name))
    print(set(authors_name))