🏷️ Make python linters happy

backend/backend/db/base_class.py

@@ -1,34 +1,39 @@
-from collections.abc import Iterable
+from __future__ import annotations
+
+from typing import Iterable, Self, TYPE_CHECKING
 
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import declarative_base
-from typing import Iterable, Self
+from sqlalchemy.orm import DeclarativeBase
 
-Base = declarative_base()
-Base.db = None
+if TYPE_CHECKING:
+    from .db import Database
 
 
-async def base_add(cls, stops: Self | Iterable[Self]) -> bool:
-    try:
-        method = (
-            cls.db.session.add_all
-            if isinstance(stops, Iterable)
-            else cls.db.session.add
-        )
-        method(stops)
-        await cls.db.session.commit()
-    except IntegrityError as err:
-        print(err)
+class Base(DeclarativeBase):
+    db: Database | None = None
 
+    @classmethod
+    async def add(cls, stops: Self | Iterable[Self]) -> bool:
+        try:
+            if isinstance(stops, Iterable):
+                cls.db.session.add_all(stops)  # type: ignore
+            else:
+                cls.db.session.add(stops)  # type: ignore
+            await cls.db.session.commit()  # type: ignore
+        except (AttributeError, IntegrityError) as err:
+            print(err)
+            return False
 
-Base.add = classmethod(base_add)
+        return True
 
-
-async def base_get_by_id(cls, id_: int | str) -> None | Base:
-    res = await cls.db.session.execute(select(cls).where(cls.id == id_))
-    element = res.scalar_one_or_none()
-    return element
-
-
-Base.get_by_id = classmethod(base_get_by_id)
+    @classmethod
+    async def get_by_id(cls, id_: int | str) -> Self | None:
+        try:
+            stmt = select(cls).where(cls.id == id_)  # type: ignore
+            res = await cls.db.session.execute(stmt)  # type: ignore
+            element = res.scalar_one_or_none()
+        except AttributeError as err:
+            print(err)
+            element = None
+        return element