from core.models.Model import Model from dataclasses import dataclass, field from dataclasses_json import config, Exclude from pytz import country_timezones from typing import Optional _table_name: str = 'locations' _table_definition: str = """ 'id' int UNIQUE, 'code' varchar UNIQUE, 'name' varchar UNIQUE """ @dataclass class Location(Model): code: str id: Optional[int] = field( default=None, metadata=config(exclude=Exclude.ALWAYS) ) name: Optional[str] = field( default=None, metadata=config(exclude=Exclude.ALWAYS) ) time_zone: str = None available: Optional[bool] = field( default=False, metadata=config(exclude=Exclude.ALWAYS) ) def __post_init__(self): if self.time_zone is None: self.time_zone = country_timezones[self.code][0] self.available = self.exists(self.code) def is_available(self): return self.exists(self.code) @staticmethod def find_by_id(id: int): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_one('SELECT * FROM locations WHERE id = ? LIMIT 1', Location.factory, [id]) @staticmethod def find(code: str): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_one('SELECT * FROM locations WHERE code = ? LIMIT 1', Location.factory, [code]) @staticmethod def exists(code: str): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_exists('SELECT * FROM locations WHERE code = ?', [code]) @staticmethod def all(): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) return Model._query_all('SELECT * FROM locations', Location.factory) @staticmethod def truncate(): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition, drop_existing=True) @staticmethod def save_many(locations): Model._create_table_if_not_exists(table_name=_table_name, table_definition=_table_definition) Model._insert_many('INSERT INTO locations VALUES(?, ?, ?)', Location.tuple_factory, locations) @staticmethod def factory(cursor, row): local_fields = [column[0] for column in cursor.description] return Location(**{key: value for key, value in zip(local_fields, row)}) @staticmethod def tuple_factory(location): return location.id, location.code, location.name