Small class to map multiple Bloop models onto a shared table

00.bloop_multitable.py

# bloop 3.0.0
import copy
from typing import Union
from bloop import BaseModel, Column
from bloop.models import subclassof, instanceof
from bloop.types import Type

__all__ = ["Mapper"]


class Mapper:
    _base_model: BaseModel

    def __init__(self, base_model: BaseModel):
        self._base_model = base_model
    
    def prefix(self, type: Type, mapped: Union[str, Column], prefix: str) -> Column:
        new_column = self._clone_column(mapped)
        new_column.typedef = PrefixType(type, mapped.typedef, prefix)
        return new_column
    
    def static(self, type: Type, mapped: Union[str, Column], value: str) -> Column:
        if callable(value):
            default = value
        else:
            default = lambda: value
        new_column = self._clone_column(mapped)
        new_column.typedef = MappedType(type, mapped.typedef)
        new_column.default = default
        return new_column
    
    def overlay(self, type: Type, mapped: Union[str, Column]) -> Column:
        new_column = self._clone_column(mapped)
        new_column.typedef = MappedType(type, mapped.typedef)
        return new_column
    
    def _clone_column(self, column: Union[str, Column]) -> Column:
        if isinstance(column, str):
            column: Column = getattr(self._base_model, column)
        clone = copy.copy(column)
        clone._dynamo_name = column.dynamo_name
        return clone


class MappedType(Type):
    _upper: Type
    _lower: Type
    
    def __init__(self, upper_typedef: Type, lower_typedef: Type):
        super().__init__()
        upper_typedef = MappedType._instantiate(upper_typedef)
        lower_typedef = MappedType._instantiate(lower_typedef)
        
        self._upper = MappedType._instantiate(upper_typedef)
        self._lower = MappedType._instantiate(lower_typedef)

        self.backing_type = self._lower.backing_type
        self.python_type = self._upper.python_type
    
    def dynamo_dump(self, value, *, context, **kwargs):
        value = self._upper.dynamo_dump(value, context=context, **kwargs)
        return self._lower.dynamo_dump(value, context=context, **kwargs)
    
    def dynamo_load(self, value, *, context, **kwargs):
        value = self._lower.dynamo_load(value, context=context, **kwargs)
        return self._upper.dynamo_load(value, context=context, **kwargs)
    
    def _instantiate(typedef) -> Type:
        if subclassof(typedef, Type):
            typedef = typedef()
        if not instanceof(typedef, Type):
            raise TypeError(f"Expected {typedef} to be instance or subclass of Type")
        return typedef


class PrefixType(MappedType):
    def __init__(self, upper_type: Type, lower_type, prefix: str):
        super().__init__(upper_type, lower_type)
        self._prefix = prefix
    
    def dynamo_dump(self, value, *, context, **kwargs):
        value = super().dynamo_dump(value, context=context, **kwargs)
        if value:
            return f"{self._prefix}{value}"
        
    def dynamo_load(self, value, *, context, **kwargs):
        if value:
            (*_, value) = value.split(self._prefix, 1)
        return super().dynamo_load(value, context=context, **kwargs)

01.usage.py

from bloop import BaseModel, Column, String, UUID
from bloop_multitable import Mapper


class SharedBase(BaseModel):
    hk = Column(String, hash_key=True)
    rk = Column(String, range_key=True)

b = SharedBase
m = Mapper(SharedBase)


class TenantUser(SharedBase):
    tenant = m.prefix(UUID, b.hk, "userByTenant::")
    id = m.overlay(UUID, b.rk)

    
# USAGE
from uuid import uuid4
from bloop import Engine
from bloop.models import unpack_from_dynamodb
from bloop.util import dump_key

engine = Engine()
engine.bind(b, skip_table_setup=True)

u = TenantUser(tenant=uuid4(), id=uuid4())
wire = dump_key(engine, u)

same = unpack_from_dynamodb(
    attrs=wire,
    expected=TenantUser.Meta.columns,
    model=TenantUser,
    engine=engine)

print(f"on the wire: {wire}")
print(f"user: {u}")
print(f"same: {same}")