mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 10:51:36 -04:00
298 lines
10 KiB
Python
298 lines
10 KiB
Python
import pathlib
|
|
from dataclasses import field, dataclass
|
|
from decimal import Decimal
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import List, Tuple
|
|
|
|
from pint import Quantity
|
|
|
|
from wrolpi.common import logger, Base, ConfigFile
|
|
from wrolpi.db import get_db_session
|
|
from wrolpi.events import Events
|
|
from .errors import NoInventories
|
|
from .inventory import get_items, get_inventories, get_unit_registry
|
|
from .models import Inventory, Item
|
|
|
|
MY_DIR: Path = Path(__file__).parent
|
|
|
|
logger = logger.getChild(__name__)
|
|
|
|
|
|
def sum_by_key(items: List[Base], key: callable):
|
|
"""
|
|
Sum the total size of each item by the provided key function. Returns a dict containing the key, and the total_size
|
|
for that key. Combine total of like units. This means ounces and pounds will be in the same total.
|
|
"""
|
|
unit_registry = get_unit_registry()
|
|
|
|
summed = dict()
|
|
for item in items:
|
|
k = key(item)
|
|
item_size, count, unit = item.item_size, item.count, unit_registry(item.unit)
|
|
key_dim = (k, unit.dimensionality)
|
|
|
|
total_size = item_size * count * unit
|
|
|
|
try:
|
|
summed[key_dim] += total_size
|
|
except KeyError:
|
|
summed[key_dim] = total_size
|
|
|
|
summed = {k: compact_unit(v) for k, v in summed.items()}
|
|
return summed
|
|
|
|
|
|
def get_inventory_by_keys(keys: Tuple, inventory_id: int):
|
|
items = get_items(inventory_id)
|
|
|
|
summed = sum_by_key(items, lambda i: tuple(getattr(i, k) or '' for k in keys))
|
|
|
|
inventory = []
|
|
for key, quantity in sorted(summed.items(), key=lambda i: i[0][0]):
|
|
quantity = cleanup_quantity(quantity)
|
|
total_size, units = quantity.to_tuple()
|
|
unit = units[0][0]
|
|
|
|
d = dict(total_size=total_size, unit=unit)
|
|
d.update(dict(zip(keys, key[0])))
|
|
inventory.append(d)
|
|
|
|
return inventory
|
|
|
|
|
|
get_inventory_by_category = partial(get_inventory_by_keys, ('category',))
|
|
get_inventory_by_subcategory = partial(get_inventory_by_keys, ('category', 'subcategory'))
|
|
get_inventory_by_name = partial(get_inventory_by_keys, ('brand', 'name'))
|
|
INVENTORY_UNITS: dict = None
|
|
UNIT_PRECISION = 5
|
|
|
|
|
|
def initialize_inventory_units() -> dict:
|
|
global INVENTORY_UNITS
|
|
if INVENTORY_UNITS:
|
|
return INVENTORY_UNITS
|
|
|
|
unit_registry = get_unit_registry()
|
|
INVENTORY_UNITS = {
|
|
('ounce', 1): (16, unit_registry.pound),
|
|
('pound', 1): (2000, unit_registry.ton),
|
|
}
|
|
|
|
|
|
def compact_unit(quantity: Quantity) -> Quantity:
|
|
"""
|
|
Convert a Quantity to it's more readable format. Such as 2000 pounds to 1 ton.
|
|
"""
|
|
initialize_inventory_units()
|
|
|
|
number, (units,) = quantity.to_tuple()
|
|
next_unit = INVENTORY_UNITS.get(units)
|
|
if not next_unit:
|
|
# No units after this one, return as is.
|
|
return round(quantity, UNIT_PRECISION)
|
|
|
|
max_quantity, next_unit = next_unit
|
|
if number >= max_quantity:
|
|
# Number is too high for this unit, move to the next unit.
|
|
return compact_unit(quantity.to(next_unit))
|
|
|
|
# No units were necessary, return as is.
|
|
quantity = round(quantity, UNIT_PRECISION)
|
|
return quantity
|
|
|
|
|
|
def quantity_to_tuple(quantity: Quantity) -> Tuple[Decimal, Quantity]:
|
|
unit_registry = get_unit_registry()
|
|
decimal, (units,) = quantity.to_tuple()
|
|
unit, _ = units
|
|
quantity, unit = round(decimal, UNIT_PRECISION), unit_registry(unit)
|
|
return quantity, unit
|
|
|
|
|
|
def cleanup_quantity(quantity: Quantity) -> Quantity:
|
|
"""Remove trailing zeros from a Quantity."""
|
|
num, unit = quantity_to_tuple(quantity)
|
|
num = round(num, UNIT_PRECISION)
|
|
num = str(num).rstrip('0').rstrip('.')
|
|
return Decimal(num) * unit
|
|
|
|
|
|
@dataclass
|
|
class InventoriesConfigValidator:
|
|
version: int = None
|
|
inventories: list[str] = field(default_factory=list)
|
|
|
|
|
|
class InventoriesConfig(ConfigFile):
|
|
file_name = 'inventories.yaml'
|
|
default_config = dict(
|
|
inventories=[],
|
|
version=0,
|
|
)
|
|
validator = InventoriesConfigValidator
|
|
|
|
@property
|
|
def inventories(self) -> dict:
|
|
return self._config['inventories']
|
|
|
|
@inventories.setter
|
|
def inventories(self, value: dict):
|
|
self.update({'inventories': value})
|
|
|
|
def dump_config(self, file: pathlib.Path = None, send_events=False, overwrite=False):
|
|
inventories = []
|
|
with get_db_session() as session:
|
|
for inventory in get_inventories(session):
|
|
inventories.append(inventory.dict())
|
|
|
|
if not inventories:
|
|
raise NoInventories('No Inventories are in the database!')
|
|
|
|
self._config['inventories'] = inventories
|
|
super().dump_config(file, send_events, overwrite)
|
|
|
|
def import_config(self, file: pathlib.Path = None, send_events=False):
|
|
file = file or self.get_file()
|
|
|
|
# If config file doesn't exist, mark import as successful (nothing to import)
|
|
if not file.is_file():
|
|
logger.info('No inventories config file, skipping import')
|
|
self.successful_import = True
|
|
return
|
|
|
|
super().import_config(file)
|
|
|
|
inventories_data = self.inventories
|
|
# Empty inventories list = never delete DB records
|
|
if not inventories_data:
|
|
logger.info(f'No inventories in config, preserving existing DB inventories')
|
|
self.successful_import = True
|
|
return
|
|
|
|
try:
|
|
with get_db_session(commit=True) as session:
|
|
db_inventories = session.query(Inventory).all()
|
|
db_inventories_by_name = {i.name: i for i in db_inventories}
|
|
config_inventories = [i['name'] for i in inventories_data]
|
|
for inventory_dict in inventories_data:
|
|
if inventory_dict['name'] not in db_inventories_by_name:
|
|
inventory = Inventory(
|
|
name=inventory_dict['name'],
|
|
created_at=inventory_dict['created_at'],
|
|
deleted_at=inventory_dict['deleted_at'],
|
|
)
|
|
session.add(inventory)
|
|
# Get the Inventory from the DB so we can use its ID.
|
|
inventory.flush()
|
|
db_inventories_by_name[inventory.name] = inventory
|
|
|
|
# Delete any inventories no longer in the config.
|
|
for name, inventory in db_inventories_by_name.items():
|
|
if name not in config_inventories:
|
|
session.delete(inventory)
|
|
|
|
for inventory_dict in inventories_data:
|
|
# Items are complex and do not have useful primary keys. Delete all items, create new items.
|
|
session.query(Item).filter(Item.inventory_id == inventory_dict['id']).delete()
|
|
inventory = db_inventories_by_name[inventory_dict['name']]
|
|
for item in inventory_dict['items']:
|
|
del item['inventory_id']
|
|
del item['id']
|
|
item = Item(inventory=inventory, inventory_id=inventory.id, **item)
|
|
session.add(item)
|
|
item.flush()
|
|
|
|
self.successful_import = True
|
|
except Exception as e:
|
|
self.successful_import = False
|
|
message = f'Failed to import {self.file_name}'
|
|
logger.error(message, exc_info=e)
|
|
if send_events:
|
|
Events.send_config_import_failed(message)
|
|
raise
|
|
|
|
def preview_backup_import(self, backup_date: str, mode: str) -> dict:
|
|
backup_file = self._get_backup_file(backup_date)
|
|
backup_data = self.read_config_file(backup_file)
|
|
current_data = self.read_config_file() if self.get_file().is_file() else dict(inventories=[], version=0)
|
|
|
|
backup_inventories = backup_data.get('inventories', [])
|
|
current_inventories = current_data.get('inventories', [])
|
|
|
|
current_names = {inv.get('name') for inv in current_inventories if inv.get('name')}
|
|
backup_names = {inv.get('name') for inv in backup_inventories if inv.get('name')}
|
|
|
|
add = []
|
|
remove = []
|
|
unchanged = 0
|
|
|
|
for inv in backup_inventories:
|
|
name = inv.get('name')
|
|
if name and name not in current_names:
|
|
add.append(dict(type='inventory', name=name))
|
|
elif name:
|
|
unchanged += 1
|
|
if mode == 'overwrite':
|
|
for name in sorted(current_names - backup_names):
|
|
remove.append(dict(type='inventory', name=name))
|
|
|
|
return dict(mode=mode, add=add, remove=remove, unchanged=unchanged)
|
|
|
|
def import_backup(self, backup_date: str, mode: str, send_events: bool = False):
|
|
self._preserve_current_config()
|
|
import shutil
|
|
backup_file = self._get_backup_file(backup_date)
|
|
config_file = self.get_file()
|
|
|
|
if mode == 'overwrite':
|
|
config_file.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(backup_file, config_file)
|
|
elif mode == 'merge':
|
|
backup_data = self.read_config_file(backup_file)
|
|
current_data = self.read_config_file() if config_file.is_file() else dict(inventories=[], version=0)
|
|
|
|
current_names = {inv.get('name') for inv in current_data.get('inventories', []) if inv.get('name')}
|
|
merged_inventories = list(current_data.get('inventories', []))
|
|
for inv in backup_data.get('inventories', []):
|
|
if inv.get('name') and inv['name'] not in current_names:
|
|
merged_inventories.append(inv)
|
|
|
|
merged_data = dict(
|
|
inventories=merged_inventories,
|
|
version=current_data.get('version', 0) + 1,
|
|
)
|
|
config_file.parent.mkdir(parents=True, exist_ok=True)
|
|
self.write_config_data(merged_data, config_file)
|
|
|
|
self.import_config(send_events=send_events)
|
|
|
|
|
|
INVENTORIES_CONFIG: InventoriesConfig = InventoriesConfig()
|
|
TEST_INVENTORIES_CONFIG: InventoriesConfig = None
|
|
|
|
|
|
def get_inventories_config():
|
|
global TEST_INVENTORIES_CONFIG
|
|
if isinstance(TEST_INVENTORIES_CONFIG, ConfigFile):
|
|
return TEST_INVENTORIES_CONFIG
|
|
|
|
global INVENTORIES_CONFIG
|
|
return INVENTORIES_CONFIG
|
|
|
|
|
|
def set_test_inventories_config(enabled: bool):
|
|
global TEST_INVENTORIES_CONFIG
|
|
if enabled:
|
|
TEST_INVENTORIES_CONFIG = InventoriesConfig()
|
|
else:
|
|
TEST_INVENTORIES_CONFIG = None
|
|
|
|
|
|
def save_inventories_config():
|
|
"""Write all inventories and their respective items to a WROLPi Config file."""
|
|
get_inventories_config().background_dump.activate_switch()
|
|
|
|
|
|
def import_inventories_config():
|
|
get_inventories_config().import_config()
|