TelegramBot/classes/location.py

from dataclasses import dataclass
from typing import Union

from bson import ObjectId
from pytz import timezone as pytz_timezone
from pytz.tzinfo import BaseTzInfo, DstTzInfo

from classes.point import Point
from modules.database_api import col_locations


@dataclass
class Location:
    __slots__ = (
        "_id",
        "id",
        "name",
        "location",
        "country",
        "timezone",
    )

    _id: ObjectId
    id: int
    name: str
    location: Point
    country: int
    timezone: Union[BaseTzInfo, DstTzInfo]

    @classmethod
    async def get(cls, id: int):
        db_entry = await col_locations.find_one({"id": id})

        if db_entry is None:
            raise ValueError(f"No location with ID {id} found.")

        db_entry["location"] = Point(*db_entry["location"])  # type: ignore
        db_entry["timezone"] = pytz_timezone(db_entry["timezone"])  # type: ignore

        return cls(**db_entry)

    @classmethod
    async def find(cls, name: str):
        db_entry = await col_locations.find_one({"name": {"$regex": name}})

        if db_entry is None:
            raise ValueError(f"No location with name {name} found.")

        db_entry["location"] = Point(*db_entry["location"])  # type: ignore
        db_entry["timezone"] = pytz_timezone(db_entry["timezone"])  # type: ignore

        return cls(**db_entry)

    @classmethod
    async def nearby(cls, lat: float, lon: float):
        db_entry = await col_locations.find_one({"location": {"$near": [lon, lat]}})

        if db_entry is None:
            raise ValueError(f"No location near {lat}, {lon} found.")

        db_entry["location"] = Point(*db_entry["location"])  # type: ignore
        db_entry["timezone"] = pytz_timezone(db_entry["timezone"])  # type: ignore

        return cls(**db_entry)