from __future__ import annotations
from abc import ABC, abstractmethod
from .segment import BaseSegment, Segment
from ezycore.models import M, Model
from ezycore.drivers import Driver
from typing import Any, List, Tuple, Type, Dict, Iterable, Union, Optional
[docs]class BaseManager(ABC):
"""
Standard manager class, all managers should inherit this class.
Parameters
----------
locations: List[Union[str, BaseSegment]]
List of locations to manage
models: Dict[:class:`str`, :class:`Model`]
Models to intialise segments with, if segment already has model, model is overwritten
location_data: Dict[:class:`str`, Dict[:class:`str`, Any]]
Kwargs for defining segment if segment doesn't already exist. Meaning its being passed by string
"""
@staticmethod
@abstractmethod
def __seg_cls():
return BaseSegment
def __init__(
self,
locations: List[Union[str, BaseSegment]],
models: Dict[str, Type[Model]] = {},
location_data: Dict[str, Dict[str, Any]] = dict()
) -> None:
self.__locations: Dict[str, BaseSegment] = dict()
self.__models: Dict[str, Type[Model]] = models
self.__index: int = 0
for v in locations:
if type(v) == str:
self.__locations[v] = self.__seg_cls()(**location_data.get(v, dict(name=v, model=self.__models.get(v))))
elif not issubclass(v, BaseSegment):
raise TypeError('Locations provided must be of type str or inherit BaseSegment class')
else:
self.__locations[v.name] = v
self.replace_model(self.__locations[v.name], self.__models.get(v.name, v.model))
self.__models[v.name] = v.model
self.__locations[v]._set_manager(self)
self._k: Tuple[str] = tuple(self.__locations)
def _modify_loc(self) -> dict:
return self.__locations
def _modify_mod(self) -> dict:
return self.__models
###########################################################################################
##
## Generic Methods
##
###########################################################################################
[docs] @abstractmethod
def populate(self, location: str, *d, data: Iterable[M] = tuple()) -> None:
""" Populate a segment with an array of data structs
Parameters
----------
location: :class:`str`
Name of segment to populate
data: Iterable[Union[:class:`dict`, :class:`Segment`]]
Array of data to add
"""
[docs] @abstractmethod
def populate_using_driver(self, location: str, driver: Driver, **driver_kwargs) -> None:
""" Populate a segment using a driver
Parameters
----------
location: :class:`str`
Name of segment to populate
driver: :class:`Driver`
Driver to use
**driver_kwargs:
Additional kwargs for :meth:`Driver.fetch`
"""
[docs] @abstractmethod
def export_segment(self, location: str, driver: Driver = None, **driver_kwargs) -> None:
""" Export a segment using a driver or any class which can handle a `export` method
Parameters
----------
location: :class:`str`
Name of segment to export
driver: :class:`Driver`
Driver to use
**driver_kwargs:
Additional kwargs for :meth:`Driver.export`
"""
###########################################################################################
##
## Segments
##
###########################################################################################
[docs] def segments(self) -> Iterable[BaseSegment]:
""" Returns all segments assigned to manager """
return iter(self.__locations.values())
[docs] def get_segment(self, segment: str, *, defer: bool = False) -> Optional[BaseSegment]:
""" Gets a saved segment, if deferred, returns `None` if segment doesn't exist.
Else raises `KeyError`.
Parameters
----------
segment: :class:`str`
Segment to get
defer: :class:`bool`
Whether to return `None` if segment not found
"""
try:
return self.__locations[segment]
except KeyError as err:
if defer:
return
raise KeyError('Segment not found') from err
[docs] @abstractmethod
def add_segment(self, segment: Union[str, BaseSegment], **kwds) -> None:
""" Add a new segment to manager,
should raise `ValueError` if segment already exists
Parameters
----------
segment: Union[:class:`str`, :class:`BaseSegment`]
Segment to add
**kwds:
If segment being added passed as string,
additional kwargs to control creation
.. note::
`model` kwarg MUST be provided
"""
[docs] @abstractmethod
def remove_segment(self, location: str, *default) -> Optional[Segment]:
""" Removes an existing segment from manager,
Should raise `ValueError` if segment doesn't exist, unless default is provided
Parameters
----------
location: :class:`str`
Name of segment to remove
*default: Any
Value to return instead of segment
"""
[docs] @abstractmethod
def update_segment(self, location: str, **data) -> BaseSegment:
""" Updates an existing segment to handle new data
Parameters
----------
location: :class:`str`
Name of segment to update
**data:
See :meth:`BaseSegment.update` for more info
"""
[docs] def replace_segment(self, segment: Union[str, BaseSegment], new_segment: Union[str, BaseSegment]) -> None:
""" Replaces/Adds a segment regardless if segment already exists,
WARNING: Inproper use may lead to unexpected errors.
* If segment provided not found, new segment is created
* If segment found, segment to be deleted and new segment is added
Parameters
----------
segment: Union[:class:`str`, :class:`BaseSegment`]
Segment to replace
new_segment: Union[:class:`str`, :class:`BaseSegment`]
Segment to use instead of existing segment
"""
try:
self.remove_segment(getattr(segment, 'name', segment))
except ValueError:
pass
self.add_segment(new_segment)
###########################################################################################
##
## Models
##
###########################################################################################
[docs] def models(self) -> Iterable[Model]:
""" Returns all models assigned to manager """
return iter(self.__models.values())
[docs] def get_model(self, location: str, *, defer: bool = False, skip_manager: bool = False) -> Optional[Model]:
""" Retrieves a model being used by a segment.
If deferred, returns `None` if model doesn't exist.
Else raises `KeyError`.
Parameters
----------
location: :class:`str`
Name of segment to retrieve model from
defer: :class:`bool`
Whether to return `None` if model not found
skip_manager: :class:`bool`
Whether to directly search segment instead of manager
"""
if skip_manager:
return self.get_segment(location).model
try:
return self.__models[location]
except KeyError as err:
if not defer:
raise KeyError('Segment model not found') from err
[docs] def replace_model(self, location: str, model: Model) -> None:
""" Adds/Overwrites model to a segment.
If model is None, then segment model is removed
Parameters
----------
location: :class:`str`
Name of segment
model: :class:`Model`
Model to bind to segment
"""
segment = self.get_segment(location)
segment.update_segment(model=model)
###########################################################################################
##
## Other Methods
##
###########################################################################################
def __repr__(self) -> str:
return f"{self.__class__.__name__}(locations={self._k})"
def __getitem__(self, segment: str) -> Segment:
try:
return self.__locations[segment]
except KeyError as err:
raise KeyError("Segment not found") from err
def __setitem__(self, segment: str, new_segment: Segment) -> None:
if not isinstance(new_segment, Segment):
raise TypeError('new_segment must of type "Segment" not "{}"'.format(new_segment.__class__.__name__))
self.replace_segment(segment, new_segment)
def __delitem__(self, segment: str) -> None:
self.remove_segment(segment)
def __contains__(self, segment: str) -> None:
assert type(segment) == str, f"Cannot compare str with {segment.__class__.__name__}"
return segment in self.__locations
def __enter__(self):
return self
def __exit__(self, *_):
return
def __iter__(self):
self.__index = 0
return self
def __next__(self):
if self.__index >= len(self._k):
raise StopIteration()
self.__index += 1
return self.__locations[self._k[self.__index - 1]]
########################################################################
## ##
## DEFAULT IMPLMENTATION ##
## ##
########################################################################
[docs]class Manager(BaseManager):
""" Default manager class implementation
Parameters
----------
locations: List[Union[str, BaseSegment]]
List of locations to manage
models: Dict[:class:`str`, :class:`Model`]
Models to intialise segments with, if segment already has model, model is overwritten
location_data: Dict[:class:`str`, Dict[:class:`str`, Any]]
Kwargs for defining segment if segment doesn't already exist. Meaning its being passed by string
"""
@staticmethod
def _BaseManager__seg_cls():
return Segment
###########################################################################################
##
## Generic Methods
##
###########################################################################################
[docs] def populate(self, location: str, *d, data: Iterable[M] = tuple()) -> None:
d = tuple(data) + d
seg = self.get_segment(location)
for loc in d:
seg.add(loc)
[docs] def populate_using_driver(self, location: str, driver: Driver, **driver_kwargs) -> None:
seg = self.get_segment(location)
if not driver_kwargs.get('model'):
driver_kwargs['model'] = seg.model
for loc in driver.fetch(location, **driver_kwargs):
seg.add(loc)
[docs] def export_segment(self, location: str, driver: Driver = None, **driver_kwargs) -> None:
seg = self.get_segment(location)
driver.export(location, seg, **driver_kwargs)
###########################################################################################
##
## Segments
##
###########################################################################################
[docs] def add_segment(self, segment: Union[str, Segment], **kwds) -> None:
seg_name = getattr(segment, 'name', segment)
_loc = self._modify_loc()
_mod = self._modify_mod()
assert type(seg_name) == str
if self.get_segment(seg_name, defer=True):
raise ValueError('Segment already exists')
if (type(segment) == str):
kwds.update(name=seg_name)
_loc[seg_name] = Segment(**kwds)
else:
_loc[seg_name] = segment
_loc[seg_name]._set_manager(self)
_mod[seg_name] = _loc[seg_name].model
self._k = tuple(_loc)
[docs] def remove_segment(self, location: str, *default) -> Optional[Segment]:
_loc = self._modify_loc()
_mod = self._modify_mod()
if location not in _loc:
if default:
return default[0]
raise ValueError('Segment not found')
rv = _loc.pop(location)
rv._del_manager()
_mod.pop(rv.name, None)
self._k = tuple(_loc)
return rv
[docs] def update_segment(self, location: str, **data) -> Segment:
segment = self.get_segment(location)
segment.update_segment(**data)