"""Role service for Atheneum.""" from collections import defaultdict from enum import Enum from typing import Optional, List, Set, Dict class Role(Enum): """User role definitions.""" ADMIN = 'ADMIN' AUDITOR = 'AUDITOR' MODERATOR = 'MODERATOR' USER = 'USER' ANONYMOUS = 'ANONYMOUS' NONE = 'NONE' class RoleTree(defaultdict): """Simple tree structure to handle hierarchy.""" def __call__(self, data: Role) -> 'RoleTree': """Handle direct calls to the tree.""" return RoleTree(self, data) # def __hash__(self): def __init__( self, parent: Optional['RoleTree'], data: Role, **kwargs: dict) -> None: """Configure a RoleTree.""" super().__init__(**kwargs) self.parent: Optional[RoleTree] = parent self.data: Role = data self.default_factory = self # type: ignore self.roles: Dict[Role, List[RoleTree]] = {data: [self]} def populate( self, children: Dict[Role, Optional[dict]]) -> List['RoleTree']: """Populate a RoleTree from a dictionary of a Role hierarchy.""" role_list: List[RoleTree] = [] for child_role in children.keys(): element = children[child_role] new_node = self(child_role) if isinstance(element, dict) and element: role_list += new_node.populate(element) self[child_role] = new_node role_list.append(new_node) for role_tree in role_list: if role_tree.data not in self.roles.keys(): self.roles[role_tree.data] = [] self.roles[role_tree.data].append(role_tree) return role_list def find_role(self, request_role: Role) -> List['RoleTree']: """Identify all instances of a role.""" try: return [role_tree for role_tree in self.roles[request_role]] except KeyError: return [] def get_parent_roles(self) -> List[Role]: """Return all the roles from self to the highest parent.""" if self.parent is not None: return [self.data] + self.parent.get_parent_roles() return [self.data] def find_roles_in_hierarchy(self, request_role: Role) -> Set[Role]: """Find a set of all roles that fall within the hierarchy.""" roles: List[Role] = [] role_trees = self.find_role(request_role) for role_tree in role_trees: roles.extend(role_tree.get_parent_roles()) return set(role for role in roles) ROLES = RoleTree(None, Role.ADMIN) ROLES.populate({ Role.MODERATOR: { Role.USER: { Role.ANONYMOUS: None } }, Role.AUDITOR: { Role.USER: None } })