sp-hydra-veil-core/core/models/Location.py

79 lines
2.6 KiB
Python

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