import json from typing import Any, Callable, Generic, Type, TypeVar, TYPE_CHECKING from urllib.error import HTTPError from flask import current_app from rebrick import lego from .exceptions import NotFoundException, ErrorException if TYPE_CHECKING: from .minifigure import BrickMinifigure from .part import BrickPart from .set import BrickSet from .socket import BrickSocket from .wish import BrickWish T = TypeVar('T', 'BrickSet', 'BrickPart', 'BrickMinifigure', 'BrickWish') # An helper around the rebrick library, autoconverting class Rebrickable(Generic[T]): method: Callable method_name: str number: str model: Type[T] socket: 'BrickSocket | None' brickset: 'BrickSet | None' minifigure: 'BrickMinifigure | None' kind: str def __init__( self, method: str, number: str, model: Type[T], /, socket: 'BrickSocket | None' = None, brickset: 'BrickSet | None' = None, minifigure: 'BrickMinifigure | None' = None ): if not hasattr(lego, method): raise ErrorException('{method} is not a valid method for the rebrick.lego module'.format( # noqa: E501 method=method, )) self.method = getattr(lego, method) self.method_name = method self.number = number self.model = model self.socket = socket self.brickset = brickset self.minifigure = minifigure if self.minifigure is not None: self.kind = 'Minifigure' else: self.kind = 'Set' # Get one element from the Rebrickable API def get(self, /) -> T: model_parameters = self.model_parameters() return self.model( **model_parameters, record=self.model.from_rebrickable( self.load(), brickset=self.brickset, ), ) # Get paginated elements from the Rebrickable API def list(self, /) -> list[T]: model_parameters = self.model_parameters() results: list[T] = [] # Bootstrap a first set of parameters parameters: dict[str, Any] | None = { 'page_size': current_app.config['REBRICKABLE_PAGE_SIZE'].value, } # Read all pages while parameters is not None: response = self.load(parameters=parameters) # Grab the results if 'results' not in response: raise ErrorException('Missing "results" field from {method} for {number}'.format( # noqa: E501 method=self.method_name, number=self.number, )) # Update the total if self.socket is not None: self.socket.total_progress(len(response['results']), add=True) # Convert to object for result in response['results']: results.append( self.model( **model_parameters, record=self.model.from_rebrickable(result), ) ) # Check for a next page if 'next' in response and response['next'] is not None: parameters['page'] = response['next'] else: parameters = None return results # Load from the API def load(self, /, parameters: dict[str, Any] = {}) -> dict[str, Any]: # Inject the API key parameters['api_key'] = current_app.config['REBRICKABLE_API_KEY'].value, # noqa: E501 try: return json.loads( self.method( self.number, **parameters, ).read() ) # HTTP errors except HTTPError as e: # Not found if e.code == 404: raise NotFoundException('{kind} {number} was not found on Rebrickable'.format( # noqa: E501 kind=self.kind, number=self.number, )) else: # Re-raise as ErrorException raise ErrorException(e) # Other errors except Exception as e: # Re-raise as ErrorException raise ErrorException(e) # Get the model parameters def model_parameters(self, /) -> dict[str, Any]: parameters: dict[str, Any] = {} # Overload with objects if self.brickset is not None: parameters['brickset'] = self.brickset if self.minifigure is not None: parameters['minifigure'] = self.minifigure return parameters